Skip to content

test: add jaxley benchmark #1896

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Draft
avik-pal wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

test: add jaxley benchmark #1896

avik-pal wants to merge 4 commits into from

Conversation

@avik-pal
Copy link
Collaborator

@avik-pal avik-pal commented Jan 7, 2026

No description provided.

@avik-pal avik-pal force-pushed the ap/neuro_benchmark branch 7 times, most recently from 94c8aa8 to 9f60d0f Compare January 7, 2026 05:32
github-actions[bot]
github-actions bot reviewed Jan 7, 2026
View reviewed changes
Copy link
Contributor

@github-actions github-actions bot left a comment •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

EnzymeJAX Benchmarks

Details
Benchmark suite Current: 9ba4982 Previous: d083e29 Ratio
actmtch / JaXPipe / cpu / Primal 0.000007481380025637918 s 0.000007391040071524912 s 1.01
actmtch / Jax / cpu / Primal 0.000007445639985235175 s 0.000007562260007034638 s 0.98
actmtch / HLOOpt / cpu / Primal 0.000007779540019328125 s 0.000011601599999266907 s 0.67
actmtch / PartOpt / cpu / Primal 0.0000073469600010867 s 0.000007577580017823493 s 0.97
actmtch / IPartOpt / cpu / Primal 0.000007591699986733147 s 0.000006997520004006219 s 1.08
actmtch / DefOpt / cpu / Primal 0.00000812714004496229 s 0.00001217241997437668 s 0.67
actmtch / IDefOpt / cpu / Primal 0.000007718300003034528 s 0.00000807944000371208 s 0.96
actmtch / JaXPipe / cpu / Forward 0.000011817159975180405 s 0.000011693580045175624 s 1.01
actmtch / Jax / cpu / Forward 0.000010585760028334336 s 0.000011192800002390868 s 0.95
actmtch / HLOOpt / cpu / Forward 0.000016300480001518736 s 0.000016192819975913154 s 1.01
actmtch / PartOpt / cpu / Forward 0.000016184699998120777 s 0.000016015620030884747 s 1.01
actmtch / IPartOpt / cpu / Forward 0.000011357099983797524 s 0.000011876220005433425 s 0.96
actmtch / DefOpt / cpu / Forward 0.00001629530002901447 s 0.000016005079969545478 s 1.02
actmtch / IDefOpt / cpu / Forward 0.000011411459954615566 s 0.000011197740032002911 s 1.02
actmtch / JaXPipe / cpu / PreRev 0.00001208606000545842 s 0.000012147400002504583 s 0.99
actmtch / JaXPipe / cpu / PostRev 0.000011714360034602578 s 0.000011047780008084374 s 1.06
actmtch / JaXPipe / cpu / BothRev 0.00001252790000762616 s 0.00001310317994466459 s 0.96
actmtch / Jax / cpu / BothRev 0.000011125759965580073 s 0.00001115773999117664 s 1.00
actmtch / HLOOpt / cpu / PreRev 0.000012365280008452828 s 0.00001213313996231591 s 1.02
actmtch / HLOOpt / cpu / PostRev 0.000016949419969023438 s 0.000016364479988624226 s 1.04
actmtch / HLOOpt / cpu / BothRev 0.00001506540003902046 s 0.00001449323996894236 s 1.04
actmtch / PartOpt / cpu / PreRev 0.000012362820025373366 s 0.000012657140041483216 s 0.98
actmtch / PartOpt / cpu / PostRev 0.000011064160016758253 s 0.00001099301995964197 s 1.01
actmtch / PartOpt / cpu / BothRev 0.00001282240007640212 s 0.000012769560016749891 s 1.00
actmtch / IPartOpt / cpu / PreRev 0.000013014240039410652 s 0.000012451840038920636 s 1.05
actmtch / IPartOpt / cpu / PostRev 0.000010856720009542187 s 0.000010973939997711567 s 0.99
actmtch / IPartOpt / cpu / BothRev 0.00001256966001164983 s 0.000012715380007648491 s 0.99
actmtch / DefOpt / cpu / PreRev 0.000012629759967239806 s 0.000012457099992388977 s 1.01
actmtch / DefOpt / cpu / PostRev 0.00001301668002270162 s 0.000012909559973195429 s 1.01
actmtch / DefOpt / cpu / BothRev 0.00001220807994286588 s 0.00001202078001369955 s 1.02
actmtch / IDefOpt / cpu / PreRev 0.000012361040035102633 s 0.00001288414000555349 s 0.96
actmtch / IDefOpt / cpu / PostRev 0.00001288267997551884 s 0.00001246106001417502 s 1.03
actmtch / IDefOpt / cpu / BothRev 0.000012716360006379543 s 0.00001241245998244267 s 1.02
actmtch / JaXPipe / cuda / Primal 0.000002016 s 0.000002016 s 1
actmtch / Jax / cuda / Primal 0.000002047 s 0.000002016 s 1.02
actmtch / HLOOpt / cuda / Primal 0.000002016 s 0.000002016 s 1
actmtch / PartOpt / cuda / Primal 0.000002016 s 0.000002016 s 1
actmtch / IPartOpt / cuda / Primal 0.000002047 s 0.000002015 s 1.02
actmtch / DefOpt / cuda / Primal 0.000002016 s 0.000002015 s 1.00
actmtch / IDefOpt / cuda / Primal 0.000002016 s 0.000002016 s 1
actmtch / JaXPipe / cuda / Forward 0.0000104 s 0.000010943 s 0.95
actmtch / Jax / cuda / Forward 0.000010016 s 0.000011072 s 0.90
actmtch / HLOOpt / cuda / Forward 0.000010432 s 0.000010752 s 0.97
actmtch / PartOpt / cuda / Forward 0.000010273 s 0.000010848 s 0.95
actmtch / IPartOpt / cuda / Forward 0.000010208 s 0.00001088 s 0.94
actmtch / DefOpt / cuda / Forward 0.00000992 s 0.000010784 s 0.92
actmtch / IDefOpt / cuda / Forward 0.000011488 s 0.000010688 s 1.07
actmtch / JaXPipe / cuda / PreRev 0.000011520000000000002 s 0.000010912 s 1.06
actmtch / JaXPipe / cuda / PostRev 0.000011296 s 0.000010368 s 1.09
actmtch / JaXPipe / cuda / BothRev 0.000011904 s 0.000010465 s 1.14
actmtch / Jax / cuda / BothRev 0.000010144 s 0.000010401 s 0.98
actmtch / HLOOpt / cuda / PreRev 0.000010144 s 0.00001024 s 0.99
actmtch / HLOOpt / cuda / PostRev 0.00001008 s 0.000010433 s 0.97
actmtch / HLOOpt / cuda / BothRev 0.000010208 s 0.000010688 s 0.96
actmtch / PartOpt / cuda / PreRev 0.00000992 s 0.000010624 s 0.93
actmtch / PartOpt / cuda / PostRev 0.000010208 s 0.000010368 s 0.98
actmtch / PartOpt / cuda / BothRev 0.000010272 s 0.000010848 s 0.95
actmtch / IPartOpt / cuda / PreRev 0.000011296 s 0.000010752 s 1.05
actmtch / IPartOpt / cuda / PostRev 0.000011584 s 0.000011136 s 1.04
actmtch / IPartOpt / cuda / BothRev 0.000011520000000000002 s 0.000010624 s 1.08
actmtch / DefOpt / cuda / PreRev 0.000010336 s 0.000010976 s 0.94
actmtch / DefOpt / cuda / PostRev 0.000010208 s 0.000010496 s 0.97
actmtch / DefOpt / cuda / BothRev 0.000009984 s 0.000010304 s 0.97
actmtch / IDefOpt / cuda / PreRev 0.000010016 s 0.000010688 s 0.94
actmtch / IDefOpt / cuda / PostRev 0.000010337 s 0.000010848 s 0.95
actmtch / IDefOpt / cuda / BothRev 0.000009984 s 0.000010336 s 0.97
actmtch / JaXPipe / tpu / Primal 5.6345e-7 s 5.6395e-7 s 1.00
actmtch / Jax / tpu / Primal 6.06325e-7 s 6.068e-7 s 1.00
actmtch / HLOOpt / tpu / Primal 0.000002101875 s 0.00000210705 s 1.00
actmtch / PartOpt / tpu / Primal 6.06675e-7 s 6.070500000000001e-7 s 1.00
actmtch / IPartOpt / tpu / Primal 5.62625e-7 s 5.63075e-7 s 1.00
actmtch / DefOpt / tpu / Primal 0.0000021703 s 0.00000216025 s 1.00
actmtch / IDefOpt / tpu / Primal 0.0000020999 s 0.000002097475 s 1.00
actmtch / JaXPipe / tpu / Forward 0.000003824775 s 0.000003818775 s 1.00
actmtch / Jax / tpu / Forward 0.00000122035 s 0.0000012161 s 1.00
actmtch / HLOOpt / tpu / Forward 0.000003941824999999999 s 0.00000394715 s 1.00
actmtch / PartOpt / tpu / Forward 0.00000391985 s 0.000003909425 s 1.00
actmtch / IPartOpt / tpu / Forward 0.000003924175 s 0.000003941824999999999 s 1.00
actmtch / DefOpt / tpu / Forward 0.0000039165 s 0.000003914725 s 1.00
actmtch / IDefOpt / tpu / Forward 0.000003934925 s 0.000003934524999999999 s 1.00
actmtch / JaXPipe / tpu / PreRev 0.000003477625 s 0.0000034856 s 1.00
actmtch / JaXPipe / tpu / PostRev 0.00000165325 s 0.000001632325 s 1.01
actmtch / JaXPipe / tpu / BothRev 0.000003496075 s 0.000003464575 s 1.01
actmtch / Jax / tpu / BothRev 0.0000016427249999999998 s 0.0000016370999999999998 s 1.00
actmtch / HLOOpt / tpu / PreRev 0.00000348765 s 0.00000347195 s 1.00
actmtch / HLOOpt / tpu / PostRev 0.0000034028500000000004 s 0.0000034105750000000004 s 1.00
actmtch / HLOOpt / tpu / BothRev 0.000003476125 s 0.0000034703 s 1.00
actmtch / PartOpt / tpu / PreRev 0.000003402175 s 0.00000341135 s 1.00
actmtch / PartOpt / tpu / PostRev 0.0000015855749999999998 s 0.0000015893 s 1.00
actmtch / PartOpt / tpu / BothRev 0.000003428325 s 0.0000034122750000000003 s 1.00
actmtch / IPartOpt / tpu / PreRev 0.0000034674 s 0.000003466875 s 1.00
actmtch / IPartOpt / tpu / PostRev 0.000001650025 s 0.0000016321 s 1.01
actmtch / IPartOpt / tpu / BothRev 0.0000034763 s 0.00000348105 s 1.00
actmtch / DefOpt / tpu / PreRev 0.0000034064 s 0.00000341815 s 1.00
actmtch / DefOpt / tpu / PostRev 0.000003407925 s 0.000003435625 s 0.99
actmtch / DefOpt / tpu / BothRev 0.000003409725 s 0.0000034119500000000004 s 1.00
actmtch / IDefOpt / tpu / PreRev 0.000003469075 s 0.000003478775 s 1.00
actmtch / IDefOpt / tpu / PostRev 0.0000034223500000000004 s 0.0000034111000000000003 s 1.00
actmtch / IDefOpt / tpu / BothRev 0.0000034648250000000003 s 0.0000034694 s 1.00
actmtch / JaXPipe / cpu / Primal 0.000013373 s 0.000007391040071524912 s 1.81
actmtch / Jax / cpu / Primal 0.000013282 s 0.000007562260007034638 s 1.76
actmtch / HLOOpt / cpu / Primal 0.000013819 s 0.000011601599999266907 s 1.19
actmtch / PartOpt / cpu / Primal 0.000013151 s 0.000007577580017823493 s 1.74
actmtch / IPartOpt / cpu / Primal 0.000013131 s 0.000006997520004006219 s 1.88
actmtch / DefOpt / cpu / Primal 0.000013805 s 0.00001217241997437668 s 1.13
actmtch / IDefOpt / cpu / Primal 0.000013519 s 0.00000807944000371208 s 1.67
actmtch / JaXPipe / cpu / Forward 0.000018992 s 0.000011693580045175624 s 1.62
actmtch / Jax / cpu / Forward 0.000017561 s 0.000011192800002390868 s 1.57
actmtch / HLOOpt / cpu / Forward 0.000018918 s 0.000016192819975913154 s 1.17
actmtch / PartOpt / cpu / Forward 0.000018593 s 0.000016015620030884747 s 1.16
actmtch / IPartOpt / cpu / Forward 0.000018425 s 0.000011876220005433425 s 1.55
actmtch / DefOpt / cpu / Forward 0.000018438 s 0.000016005079969545478 s 1.15
actmtch / IDefOpt / cpu / Forward 0.000018983 s 0.000011197740032002911 s 1.70
actmtch / JaXPipe / cpu / PreRev 0.00001907 s 0.000012147400002504583 s 1.57
actmtch / JaXPipe / cpu / PostRev 0.000017454999999999998 s 0.000011047780008084374 s 1.58
actmtch / JaXPipe / cpu / BothRev 0.000019258 s 0.00001310317994466459 s 1.47
actmtch / Jax / cpu / BothRev 0.000017622 s 0.00001115773999117664 s 1.58
actmtch / HLOOpt / cpu / PreRev 0.000018547 s 0.00001213313996231591 s 1.53
actmtch / HLOOpt / cpu / PostRev 0.000019307 s 0.000016364479988624226 s 1.18
actmtch / HLOOpt / cpu / BothRev 0.000019519 s 0.00001449323996894236 s 1.35
actmtch / PartOpt / cpu / PreRev 0.000019027 s 0.000012657140041483216 s 1.50
actmtch / PartOpt / cpu / PostRev 0.000017276 s 0.00001099301995964197 s 1.57
actmtch / PartOpt / cpu / BothRev 0.00001918 s 0.000012769560016749891 s 1.50
actmtch / IPartOpt / cpu / PreRev 0.000019148 s 0.000012451840038920636 s 1.54
actmtch / IPartOpt / cpu / PostRev 0.000017651 s 0.000010973939997711567 s 1.61
actmtch / IPartOpt / cpu / BothRev 0.00001889 s 0.000012715380007648491 s 1.49
actmtch / DefOpt / cpu / PreRev 0.000018847 s 0.000012457099992388977 s 1.51
actmtch / DefOpt / cpu / PostRev 0.000019015 s 0.000012909559973195429 s 1.47
actmtch / DefOpt / cpu / BothRev 0.000019432000000000003 s 0.00001202078001369955 s 1.62
actmtch / IDefOpt / cpu / PreRev 0.000018563 s 0.00001288414000555349 s 1.44
actmtch / IDefOpt / cpu / PostRev 0.000019168 s 0.00001246106001417502 s 1.54
actmtch / IDefOpt / cpu / BothRev 0.00001929 s 0.00001241245998244267 s 1.55
add_one / JaXPipe / cpu / Primal 0.000008278660043288256 s 0.000007896179986346397 s 1.05
add_one / Jax / cpu / Primal 0.000007530059992859605 s 0.000008059760020842077 s 0.93
add_one / HLOOpt / cpu / Primal 0.000010662720014806836 s 0.000011271700041106667 s 0.95
add_one / PartOpt / cpu / Primal 0.000007706939995841822 s 0.00000737778007533052 s 1.04
add_one / IPartOpt / cpu / Primal 0.000007406939985230565 s 0.000007323179988816264 s 1.01
add_one / DefOpt / cpu / Primal 0.000007663720061827917 s 0.000011454540026534232 s 0.67
add_one / IDefOpt / cpu / Primal 0.000007327700013775029 s 0.000007599880000270786 s 0.96
add_one / JaXPipe / cpu / Forward 0.000011060960014219744 s 0.000011489840044305313 s 0.96
add_one / Jax / cpu / Forward 0.000011795899981734692 s 0.00001128000002609042 s 1.05
add_one / HLOOpt / cpu / Forward 0.00001579429999765125 s 0.000015948640002534375 s 0.99
add_one / PartOpt / cpu / Forward 0.000016525819955859333 s 0.00001568542000313755 s 1.05
add_one / IPartOpt / cpu / Forward 0.000011438339997766889 s 0.000011263400028838078 s 1.02
add_one / DefOpt / cpu / Forward 0.000016111760041894742 s 0.00001617061998331337 s 1.00
add_one / IDefOpt / cpu / Forward 0.000011472059977677418 s 0.000011673080025502714 s 0.98
add_one / JaXPipe / cpu / PreRev 0.000013097599976390484 s 0.00001287380003304861 s 1.02
add_one / JaXPipe / cpu / PostRev 0.000012681139987762436 s 0.000012663920015256736 s 1.00
add_one / JaXPipe / cpu / BothRev 0.000016697860019121435 s 0.000017538880010761206 s 0.95
add_one / Jax / cpu / BothRev 0.000012668300014411216 s 0.000012782919984601904 s 0.99
add_one / HLOOpt / cpu / PreRev 0.000012617040001714484 s 0.000013180160012780106 s 0.96
add_one / HLOOpt / cpu / PostRev 0.00001701595999293204 s 0.000012923799995405717 s 1.32
add_one / HLOOpt / cpu / BothRev 0.000015062700031194254 s 0.000014851300011287096 s 1.01
add_one / PartOpt / cpu / PreRev 0.000012612879991138471 s 0.000013053980010226951 s 0.97
add_one / PartOpt / cpu / PostRev 0.000013213920001362566 s 0.00001259243996173609 s 1.05
add_one / PartOpt / cpu / BothRev 0.000013077820012767916 s 0.000013003179947190802 s 1.01
add_one / IPartOpt / cpu / PreRev 0.000017897760017149267 s 0.000015980460029823008 s 1.12
add_one / IPartOpt / cpu / PostRev 0.00001278895999348606 s 0.00001266117998966365 s 1.01
add_one / IPartOpt / cpu / BothRev 0.000012596659998962423 s 0.0000128849999873637 s 0.98
add_one / DefOpt / cpu / PreRev 0.00001285329998609086 s 0.000013108160019328352 s 0.98
add_one / DefOpt / cpu / PostRev 0.000013301020007929765 s 0.000012861699933637284 s 1.03
add_one / DefOpt / cpu / BothRev 0.000013234060015747672 s 0.000012957519984411193 s 1.02
add_one / IDefOpt / cpu / PreRev 0.000013185040006646889 s 0.000013135819945091498 s 1.00
add_one / IDefOpt / cpu / PostRev 0.000013146599958417937 s 0.000012651040069613371 s 1.04
add_one / IDefOpt / cpu / BothRev 0.000012830380010200317 s 0.000012945279995619786 s 0.99
add_one / JaXPipe / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_one / Jax / cuda / Primal 0.000001951 s 0.0000019200000000000003 s 1.02
add_one / HLOOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_one / PartOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_one / IPartOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_one / DefOpt / cuda / Primal 0.000001951 s 0.0000019200000000000003 s 1.02
add_one / IDefOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_one / JaXPipe / cuda / Forward 0.000010529 s 0.000010624 s 0.99
add_one / Jax / cuda / Forward 0.000010016 s 0.000010688 s 0.94
add_one / HLOOpt / cuda / Forward 0.000010144 s 0.000010496 s 0.97
add_one / PartOpt / cuda / Forward 0.000010208 s 0.00001056 s 0.97
add_one / IPartOpt / cuda / Forward 0.000010272 s 0.000010688 s 0.96
add_one / DefOpt / cuda / Forward 0.0000104 s 0.000010656 s 0.98
add_one / IDefOpt / cuda / Forward 0.00001024 s 0.000010848 s 0.94
add_one / JaXPipe / cuda / PreRev 0.000026272 s 0.000025632 s 1.02
add_one / JaXPipe / cuda / PostRev 0.000026336 s 0.000026272 s 1.00
add_one / JaXPipe / cuda / BothRev 0.000025696 s 0.000025952 s 0.99
add_one / Jax / cuda / BothRev 0.0000256 s 0.00002528 s 1.01
add_one / HLOOpt / cuda / PreRev 0.000025985000000000003 s 0.000025825 s 1.01
add_one / HLOOpt / cuda / PostRev 0.000025504 s 0.000025024 s 1.02
add_one / HLOOpt / cuda / BothRev 0.000026336 s 0.000026144 s 1.01
add_one / PartOpt / cuda / PreRev 0.00002576 s 0.000026017 s 0.99
add_one / PartOpt / cuda / PostRev 0.000026208 s 0.000026144 s 1.00
add_one / PartOpt / cuda / BothRev 0.000026529 s 0.00002544 s 1.04
add_one / IPartOpt / cuda / PreRev 0.000026016 s 0.00002592 s 1.00
add_one / IPartOpt / cuda / PostRev 0.000025888 s 0.000026016 s 1.00
add_one / IPartOpt / cuda / BothRev 0.000030048 s 0.000026144 s 1.15
add_one / DefOpt / cuda / PreRev 0.000026752 s 0.000025792 s 1.04
add_one / DefOpt / cuda / PostRev 0.000025376 s 0.00002624 s 0.97
add_one / DefOpt / cuda / BothRev 0.000025664 s 0.000025568 s 1.00
add_one / IDefOpt / cuda / PreRev 0.00002544 s 0.00002576 s 0.99
add_one / IDefOpt / cuda / PostRev 0.000025984 s 0.000025632 s 1.01
add_one / IDefOpt / cuda / BothRev 0.000026272 s 0.000025921 s 1.01
add_one / JaXPipe / tpu / Primal 0.0000014216 s 0.000001430875 s 0.99
add_one / Jax / tpu / Primal 0.000001406975 s 0.000001400375 s 1.00
add_one / HLOOpt / tpu / Primal 0.00000142385 s 0.000001429925 s 1.00
add_one / PartOpt / tpu / Primal 0.000001408575 s 0.000001403325 s 1.00
add_one / IPartOpt / tpu / Primal 0.00000142625 s 0.0000014275499999999998 s 1.00
add_one / DefOpt / tpu / Primal 0.000001413075 s 0.00000139915 s 1.01
add_one / IDefOpt / tpu / Primal 0.000001436175 s 0.0000014233000000000002 s 1.01
add_one / JaXPipe / tpu / Forward 0.000001847 s 0.000001854225 s 1.00
add_one / Jax / tpu / Forward 0.0000018357 s 0.000001834425 s 1.00
add_one / HLOOpt / tpu / Forward 0.000001847875 s 0.000001844875 s 1.00
add_one / PartOpt / tpu / Forward 0.00000185505 s 0.00000183895 s 1.01
add_one / IPartOpt / tpu / Forward 0.00000185055 s 0.000001848975 s 1.00
add_one / DefOpt / tpu / Forward 0.00000183725 s 0.000001842675 s 1.00
add_one / IDefOpt / tpu / Forward 0.000001861625 s 0.00000185815 s 1.00
add_one / JaXPipe / tpu / PreRev 0.000002233125 s 0.000002238925 s 1.00
add_one / JaXPipe / tpu / PostRev 0.00000223475 s 0.0000022499 s 0.99
add_one / JaXPipe / tpu / BothRev 0.00000223495 s 0.0000022318 s 1.00
add_one / Jax / tpu / BothRev 0.000002240075 s 0.00000223495 s 1.00
add_one / HLOOpt / tpu / PreRev 0.00000223605 s 0.000002232875 s 1.00
add_one / HLOOpt / tpu / PostRev 0.000002246125 s 0.00000223685 s 1.00
add_one / HLOOpt / tpu / BothRev 0.00000223955 s 0.000002234575 s 1.00
add_one / PartOpt / tpu / PreRev 0.000002242825 s 0.000002237975 s 1.00
add_one / PartOpt / tpu / PostRev 0.0000022455 s 0.00000223315 s 1.01
add_one / PartOpt / tpu / BothRev 0.000002244675 s 0.00000224775 s 1.00
add_one / IPartOpt / tpu / PreRev 0.0000022404 s 0.000002233025 s 1.00
add_one / IPartOpt / tpu / PostRev 0.00000224025 s 0.000002241175 s 1.00
add_one / IPartOpt / tpu / BothRev 0.0000022463 s 0.00000223925 s 1.00
add_one / DefOpt / tpu / PreRev 0.0000022418250000000003 s 0.000002235275 s 1.00
add_one / DefOpt / tpu / PostRev 0.00000223865 s 0.0000022341000000000003 s 1.00
add_one / DefOpt / tpu / BothRev 0.000002240025 s 0.00000223715 s 1.00
add_one / IDefOpt / tpu / PreRev 0.00000223345 s 0.0000022430000000000004 s 1.00
add_one / IDefOpt / tpu / PostRev 0.00000223845 s 0.000002242425 s 1.00
add_one / IDefOpt / tpu / BothRev 0.0000022357 s 0.000002245075 s 1.00
add_one / JaXPipe / cpu / Primal 0.000013216 s 0.000007896179986346397 s 1.67
add_one / Jax / cpu / Primal 0.000012833 s 0.000008059760020842077 s 1.59
add_one / HLOOpt / cpu / Primal 0.000012469 s 0.000011271700041106667 s 1.11
add_one / PartOpt / cpu / Primal 0.000020264 s 0.00000737778007533052 s 2.75
add_one / IPartOpt / cpu / Primal 0.000012553 s 0.000007323179988816264 s 1.71
add_one / DefOpt / cpu / Primal 0.000012727 s 0.000011454540026534232 s 1.11
add_one / IDefOpt / cpu / Primal 0.000012413 s 0.000007599880000270786 s 1.63
add_one / JaXPipe / cpu / Forward 0.000017514 s 0.000011489840044305313 s 1.52
add_one / Jax / cpu / Forward 0.000017017 s 0.00001128000002609042 s 1.51
add_one / HLOOpt / cpu / Forward 0.000017046 s 0.000015948640002534375 s 1.07
add_one / PartOpt / cpu / Forward 0.000017222000000000002 s 0.00001568542000313755 s 1.10
add_one / IPartOpt / cpu / Forward 0.000017072000000000002 s 0.000011263400028838078 s 1.52
add_one / DefOpt / cpu / Forward 0.000017137 s 0.00001617061998331337 s 1.06
add_one / IDefOpt / cpu / Forward 0.000017114 s 0.000011673080025502714 s 1.47
add_one / JaXPipe / cpu / PreRev 0.000019299 s 0.00001287380003304861 s 1.50
add_one / JaXPipe / cpu / PostRev 0.000019672 s 0.000012663920015256736 s 1.55
add_one / JaXPipe / cpu / BothRev 0.000019584 s 0.000017538880010761206 s 1.12
add_one / Jax / cpu / BothRev 0.000019235 s 0.000012782919984601904 s 1.50
add_one / HLOOpt / cpu / PreRev 0.000019438 s 0.000013180160012780106 s 1.47
add_one / HLOOpt / cpu / PostRev 0.000019693 s 0.000012923799995405717 s 1.52
add_one / HLOOpt / cpu / BothRev 0.000019015 s 0.000014851300011287096 s 1.28
add_one / PartOpt / cpu / PreRev 0.000019472 s 0.000013053980010226951 s 1.49
add_one / PartOpt / cpu / PostRev 0.000019931 s 0.00001259243996173609 s 1.58
add_one / PartOpt / cpu / BothRev 0.000019652 s 0.000013003179947190802 s 1.51
add_one / IPartOpt / cpu / PreRev 0.000019391 s 0.000015980460029823008 s 1.21
add_one / IPartOpt / cpu / PostRev 0.000019575 s 0.00001266117998966365 s 1.55
add_one / IPartOpt / cpu / BothRev 0.000019706 s 0.0000128849999873637 s 1.53
add_one / DefOpt / cpu / PreRev 0.000019224 s 0.000013108160019328352 s 1.47
add_one / DefOpt / cpu / PostRev 0.000019487 s 0.000012861699933637284 s 1.52
add_one / DefOpt / cpu / BothRev 0.000019284 s 0.000012957519984411193 s 1.49
add_one / IDefOpt / cpu / PreRev 0.000019514 s 0.000013135819945091498 s 1.49
add_one / IDefOpt / cpu / PostRev 0.000019335 s 0.000012651040069613371 s 1.53
add_one / IDefOpt / cpu / BothRev 0.000019414 s 0.000012945279995619786 s 1.50
add_two / JaXPipe / cpu / Primal 0.000008340199974554707 s 0.000008078299997578142 s 1.03
add_two / Jax / cpu / Primal 0.00000768115999562724 s 0.000007480699987354456 s 1.03
add_two / HLOOpt / cpu / Primal 0.000011363359953975304 s 0.000011705360002451924 s 0.97
add_two / PartOpt / cpu / Primal 0.000007584940003653174 s 0.000007847820043025422 s 0.97
add_two / IPartOpt / cpu / Primal 0.00000764456000069913 s 0.000008033739977690856 s 0.95
add_two / DefOpt / cpu / Primal 0.00001147496002886328 s 0.000011882460030392397 s 0.97
add_two / IDefOpt / cpu / Primal 0.000007544319978478597 s 0.000007653040029254043 s 0.99
add_two / JaXPipe / cpu / Forward 0.000011109940005553654 s 0.00001155064003796724 s 0.96
add_two / Jax / cpu / Forward 0.00001148011998338916 s 0.00001183109997327847 s 0.97
add_two / HLOOpt / cpu / Forward 0.00001587730002029275 s 0.000016174619977391556 s 0.98
add_two / PartOpt / cpu / Forward 0.000015894939988356782 s 0.00001656830003412324 s 0.96
add_two / IPartOpt / cpu / Forward 0.00001121786000112479 s 0.000011428399993747009 s 0.98
add_two / DefOpt / cpu / Forward 0.000016200079980990266 s 0.000011539719998836518 s 1.40
add_two / IDefOpt / cpu / Forward 0.000011970940013270593 s 0.00001197085999592673 s 1.00
add_two / JaXPipe / cpu / PreRev 0.000015492679995077197 s 0.000015447079995283275 s 1.00
add_two / JaXPipe / cpu / PostRev 0.000015630359994247555 s 0.000015476439966732868 s 1.01
add_two / JaXPipe / cpu / BothRev 0.000015730199975223514 s 0.00001541200000247045 s 1.02
add_two / Jax / cpu / BothRev 0.00001541653999083792 s 0.00001537471996016393 s 1.00
add_two / HLOOpt / cpu / PreRev 0.000015806819956196705 s 0.00001551939993078122 s 1.02
add_two / HLOOpt / cpu / PostRev 0.000016249220043391687 s 0.00001580970001668902 s 1.03
add_two / HLOOpt / cpu / BothRev 0.000016473499990752315 s 0.000017628919977141777 s 0.93
add_two / PartOpt / cpu / PreRev 0.000015034579982966534 s 0.000015835540016269078 s 0.95
add_two / PartOpt / cpu / PostRev 0.00001548687999274989 s 0.000015323180014092942 s 1.01
add_two / PartOpt / cpu / BothRev 0.00001496499999120715 s 0.000015682500015827828 s 0.95
add_two / IPartOpt / cpu / PreRev 0.00001524163999420125 s 0.000015858660008234436 s 0.96
add_two / IPartOpt / cpu / PostRev 0.000015500240024266532 s 0.000015334940017055487 s 1.01
add_two / IPartOpt / cpu / BothRev 0.00001519636002740299 s 0.000015353399994637585 s 0.99
add_two / DefOpt / cpu / PreRev 0.000015164499973252533 s 0.000015210360006676635 s 1.00
add_two / DefOpt / cpu / PostRev 0.000015764880017741235 s 0.00001574001999870234 s 1.00
add_two / DefOpt / cpu / BothRev 0.000015876339966780505 s 0.000015454039994438064 s 1.03
add_two / IDefOpt / cpu / PreRev 0.000016639499999655527 s 0.000015190960048130363 s 1.10
add_two / IDefOpt / cpu / PostRev 0.000015793959983056992 s 0.000016048160032369195 s 0.98
add_two / IDefOpt / cpu / BothRev 0.00001582380000400008 s 0.000015832279987080257 s 1.00
add_two / JaXPipe / cuda / Primal 0.000001951 s 0.0000019200000000000003 s 1.02
add_two / Jax / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_two / HLOOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_two / PartOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_two / IPartOpt / cuda / Primal 0.000001951 s 0.000001888 s 1.03
add_two / DefOpt / cuda / Primal 0.0000019200000000000003 s 0.000001889 s 1.02
add_two / IDefOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_two / JaXPipe / cuda / Forward 0.000008832 s 0.000010272 s 0.86
add_two / Jax / cuda / Forward 0.000010016 s 0.000010464 s 0.96
add_two / HLOOpt / cuda / Forward 0.00000992 s 0.000010367 s 0.96
add_two / PartOpt / cuda / Forward 0.00001008 s 0.00000976 s 1.03
add_two / IPartOpt / cuda / Forward 0.000010175 s 0.000010368 s 0.98
add_two / DefOpt / cuda / Forward 0.000010176 s 0.000009664 s 1.05
add_two / IDefOpt / cuda / Forward 0.000010113 s 0.000009952 s 1.02
add_two / JaXPipe / cuda / PreRev 0.000033119999999999995 s 0.00003344 s 0.99
add_two / JaXPipe / cuda / PostRev 0.000033632 s 0.000034049000000000006 s 0.99
add_two / JaXPipe / cuda / BothRev 0.000033184 s 0.000034464 s 0.96
add_two / Jax / cuda / BothRev 0.000033184 s 0.000032769 s 1.01
add_two / HLOOpt / cuda / PreRev 0.000033377 s 0.000033312 s 1.00
add_two / HLOOpt / cuda / PostRev 0.000032864 s 0.000033119999999999995 s 0.99
add_two / HLOOpt / cuda / BothRev 0.000032384 s 0.000033856 s 0.96
add_two / PartOpt / cuda / PreRev 0.000033759999999999995 s 0.000033568 s 1.01
add_two / PartOpt / cuda / PostRev 0.000032832 s 0.000033569 s 0.98
add_two / PartOpt / cuda / BothRev 0.000032608 s 0.000032800000000000004 s 0.99
add_two / IPartOpt / cuda / PreRev 0.000032256 s 0.000032704 s 0.99
add_two / IPartOpt / cuda / PostRev 0.000032704 s 0.00003328 s 0.98
add_two / IPartOpt / cuda / BothRev 0.000032352 s 0.000032673000000000004 s 0.99
add_two / DefOpt / cuda / PreRev 0.000032417 s 0.000033184 s 0.98
add_two / DefOpt / cuda / PostRev 0.000033057000000000006 s 0.0000336 s 0.98
add_two / DefOpt / cuda / BothRev 0.000032896000000000005 s 0.000034336 s 0.96
add_two / IDefOpt / cuda / PreRev 0.000032993 s 0.000034144000000000004 s 0.97
add_two / IDefOpt / cuda / PostRev 0.000032864 s 0.000033824 s 0.97
add_two / IDefOpt / cuda / BothRev 0.000033119999999999995 s 0.000034048 s 0.97
add_two / JaXPipe / tpu / Primal 0.0000014277500000000002 s 0.0000014336499999999998 s 1.00
add_two / Jax / tpu / Primal 0.000001480175 s 0.000001484 s 1.00
add_two / HLOOpt / tpu / Primal 0.0000014356 s 0.0000014327 s 1.00
add_two / PartOpt / tpu / Primal 0.0000014799750000000002 s 0.0000014771 s 1.00
add_two / IPartOpt / tpu / Primal 0.0000014396000000000002 s 0.0000014297499999999995 s 1.01
add_two / DefOpt / tpu / Primal 0.0000014754249999999995 s 0.0000014844250000000002 s 0.99
add_two / IDefOpt / tpu / Primal 0.000001432075 s 0.000001430075 s 1.00
add_two / JaXPipe / tpu / Forward 0.000001826975 s 0.0000018276 s 1.00
add_two / Jax / tpu / Forward 0.000001829075 s 0.000001824925 s 1.00
add_two / HLOOpt / tpu / Forward 0.000001830925 s 0.0000018248 s 1.00
add_two / PartOpt / tpu / Forward 0.0000018351 s 0.000001825575 s 1.01
add_two / IPartOpt / tpu / Forward 0.000001828875 s 0.00000182455 s 1.00
add_two / DefOpt / tpu / Forward 0.000001824 s 0.0000018259 s 1.00
add_two / IDefOpt / tpu / Forward 0.0000018258750000000003 s 0.0000018313 s 1.00
add_two / JaXPipe / tpu / PreRev 0.00000283885 s 0.000002842925 s 1.00
add_two / JaXPipe / tpu / PostRev 0.0000027566 s 0.000002745025 s 1.00
add_two / JaXPipe / tpu / BothRev 0.000002838175 s 0.000002831875 s 1.00
add_two / Jax / tpu / BothRev 0.000002744275 s 0.000002747275 s 1.00
add_two / HLOOpt / tpu / PreRev 0.0000028397750000000003 s 0.000002842225 s 1.00
add_two / HLOOpt / tpu / PostRev 0.000002756025 s 0.000002754675 s 1.00
add_two / HLOOpt / tpu / BothRev 0.0000028347499999999994 s 0.0000028454 s 1.00
add_two / PartOpt / tpu / PreRev 0.0000027615 s 0.0000027563 s 1.00
add_two / PartOpt / tpu / PostRev 0.000002832325 s 0.000002837525 s 1.00
add_two / PartOpt / tpu / BothRev 0.000002750275 s 0.000002755075 s 1.00
add_two / IPartOpt / tpu / PreRev 0.000002837825 s 0.000002825175 s 1.00
add_two / IPartOpt / tpu / PostRev 0.000002752225 s 0.00000275205 s 1.00
add_two / IPartOpt / tpu / BothRev 0.000002832525 s 0.0000028322750000000003 s 1.00
add_two / DefOpt / tpu / PreRev 0.0000027537250000000005 s 0.000002747625 s 1.00
add_two / DefOpt / tpu / PostRev 0.0000028413749999999995 s 0.000002837 s 1.00
add_two / DefOpt / tpu / BothRev 0.0000027426 s 0.000002767225 s 0.99
add_two / IDefOpt / tpu / PreRev 0.0000028459 s 0.000002828475 s 1.01
add_two / IDefOpt / tpu / PostRev 0.000002750025 s 0.0000027595 s 1.00
add_two / IDefOpt / tpu / BothRev 0.00000282715 s 0.000002844775 s 0.99
add_two / JaXPipe / cpu / Primal 0.000013325 s 0.000008078299997578142 s 1.65
add_two / Jax / cpu / Primal 0.000013096 s 0.000007480699987354456 s 1.75
add_two / HLOOpt / cpu / Primal 0.000013063 s 0.000011705360002451924 s 1.12
add_two / PartOpt / cpu / Primal 0.000013089 s 0.000007847820043025422 s 1.67
add_two / IPartOpt / cpu / Primal 0.000013286 s 0.000008033739977690856 s 1.65
add_two / DefOpt / cpu / Primal 0.000013324000000000002 s 0.000011882460030392397 s 1.12
add_two / IDefOpt / cpu / Primal 0.000013189 s 0.000007653040029254043 s 1.72
add_two / JaXPipe / cpu / Forward 0.000018193 s 0.00001155064003796724 s 1.58
add_two / Jax / cpu / Forward 0.000017673000000000002 s 0.00001183109997327847 s 1.49
add_two / HLOOpt / cpu / Forward 0.000017888999999999998 s 0.000016174619977391556 s 1.11
add_two / PartOpt / cpu / Forward 0.000017786 s 0.00001656830003412324 s 1.07
add_two / IPartOpt / cpu / Forward 0.000017542999999999998 s 0.000011428399993747009 s 1.54
add_two / DefOpt / cpu / Forward 0.000024191 s 0.000011539719998836518 s 2.10
add_two / IDefOpt / cpu / Forward 0.000017882 s 0.00001197085999592673 s 1.49
add_two / JaXPipe / cpu / PreRev 0.000022719 s 0.000015447079995283275 s 1.47
add_two / JaXPipe / cpu / PostRev 0.000023414 s 0.000015476439966732868 s 1.51
add_two / JaXPipe / cpu / BothRev 0.000023036 s 0.00001541200000247045 s 1.49
add_two / Jax / cpu / BothRev 0.000023381 s 0.00001537471996016393 s 1.52
add_two / HLOOpt / cpu / PreRev 0.000022637 s 0.00001551939993078122 s 1.46
add_two / HLOOpt / cpu / PostRev 0.000023039 s 0.00001580970001668902 s 1.46
add_two / HLOOpt / cpu / BothRev 0.000023183 s 0.000017628919977141777 s 1.32
add_two / PartOpt / cpu / PreRev 0.000022968 s 0.000015835540016269078 s 1.45
add_two / PartOpt / cpu / PostRev 0.000023314 s 0.000015323180014092942 s 1.52
add_two / PartOpt / cpu / BothRev 0.000023042 s 0.000015682500015827828 s 1.47
add_two / IPartOpt / cpu / PreRev 0.000022511 s 0.000015858660008234436 s 1.42
add_two / IPartOpt / cpu / PostRev 0.000023498 s 0.000015334940017055487 s 1.53
add_two / IPartOpt / cpu / BothRev 0.000022841 s 0.000015353399994637585 s 1.49
add_two / DefOpt / cpu / PreRev 0.000022344 s 0.000015210360006676635 s 1.47
add_two / DefOpt / cpu / PostRev 0.00002318 s 0.00001574001999870234 s 1.47
add_two / DefOpt / cpu / BothRev 0.000023239 s 0.000015454039994438064 s 1.50
add_two / IDefOpt / cpu / PreRev 0.000022717 s 0.000015190960048130363 s 1.50
add_two / IDefOpt / cpu / PostRev 0.000023102 s 0.000016048160032369195 s 1.44
add_two / IDefOpt / cpu / BothRev 0.00002284 s 0.000015832279987080257 s 1.44
cache / JaXPipe / cpu / Primal 0.000007428319995597121 s 0.000007311580002351548 s 1.02
cache / Jax / cpu / Primal 0.00000781616000494978 s 0.000008188780020645936 s 0.95
cache / HLOOpt / cpu / Primal 0.000007219939989226986 s 0.000007291099991562077 s 0.99
cache / PartOpt / cpu / Primal 0.000007324719999814988 s 0.000007267500013767858 s 1.01
cache / IPartOpt / cpu / Primal 0.000007589000051666517 s 0.000007347700038735638 s 1.03
cache / DefOpt / cpu / Primal 0.000007133259996408014 s 0.000007079820006765658 s 1.01
cache / IDefOpt / cpu / Primal 0.000007142759995986125 s 0.000007294299994100584 s 0.98
cache / JaXPipe / cpu / Forward 0.000016731140058254822 s 0.000015604340023855912 s 1.07
cache / Jax / cpu / Forward 0.000016661359995850945 s 0.000016084959997897386 s 1.04
cache / HLOOpt / cpu / Forward 0.000021697959982702745 s 0.000016539119997105445 s 1.31
cache / PartOpt / cpu / Forward 0.000021781400027975907 s 0.00002018942001996038 s 1.08
cache / IPartOpt / cpu / Forward 0.000016486719987369723 s 0.00001527753997834225 s 1.08
cache / DefOpt / cpu / Forward 0.000021046900010333046 s 0.00002051123997262039 s 1.03
cache / IDefOpt / cpu / Forward 0.000015807760028110352 s 0.00001486773998294666 s 1.06
cache / JaXPipe / cpu / PreRev 0.000018384760005574206 s 0.000016669340011503665 s 1.10
cache / JaXPipe / cpu / PostRev 0.000022782179985370025 s 0.000022038719989723177 s 1.03
cache / JaXPipe / cpu / BothRev 0.00001862460000666033 s 0.00001681013998677372 s 1.11
cache / Jax / cpu / BothRev 0.000022266340038186173 s 0.000022604220002904186 s 0.99
cache / HLOOpt / cpu / PreRev 0.00001741657998536539 s 0.000017696700033411617 s 0.98
cache / HLOOpt / cpu / PostRev 0.00001809250004043861 s 0.00002060288000393484 s 0.88
cache / HLOOpt / cpu / BothRev 0.00002053884002634732 s 0.0000200398199740448 s 1.02
cache / PartOpt / cpu / PreRev 0.000017436359948987958 s 0.000017317400006504615 s 1.01
cache / PartOpt / cpu / PostRev 0.00002227468001365196 s 0.00002674846000445541 s 0.83
cache / PartOpt / cpu / BothRev 0.00001731617996483692 s 0.00001720827998724417 s 1.01
cache / IPartOpt / cpu / PreRev 0.000017408619987691055 s 0.00001730513999973482 s 1.01
cache / IPartOpt / cpu / PostRev 0.0000228809599957458 s 0.0000223068999730458 s 1.03
cache / IPartOpt / cpu / BothRev 0.000017450180075684328 s 0.000017066479967979832 s 1.02
cache / DefOpt / cpu / PreRev 0.00001789534003364679 s 0.000017424999950890195 s 1.03
cache / DefOpt / cpu / PostRev 0.00001831939997828158 s 0.000017306580002696136 s 1.06
cache / DefOpt / cpu / BothRev 0.000016954980028458523 s 0.00001792524001757556 s 0.95
cache / IDefOpt / cpu / PreRev 0.000017833640031312824 s 0.00001806160001251556 s 0.99
cache / IDefOpt / cpu / PostRev 0.000017378939955960958 s 0.00001825745999667561 s 0.95
cache / IDefOpt / cpu / BothRev 0.000017119380017902586 s 0.00001824015998863615 s 0.94
cache / JaXPipe / cuda / Primal 0.000002335 s 0.000002272 s 1.03
cache / Jax / cuda / Primal 0.000002272 s 0.00000224 s 1.01
cache / HLOOpt / cuda / Primal 0.000002304 s 0.00000224 s 1.03
cache / PartOpt / cuda / Primal 0.000002272 s 0.00000224 s 1.01
cache / IPartOpt / cuda / Primal 0.000002273 s 0.000002208 s 1.03
cache / DefOpt / cuda / Primal 0.000002272 s 0.00000224 s 1.01
cache / IDefOpt / cuda / Primal 0.000002304 s 0.000002303 s 1.00
cache / JaXPipe / cuda / Forward 0.000002336 s 0.000002304 s 1.01
cache / Jax / cuda / Forward 0.000002337 s 0.000002272 s 1.03
cache / HLOOpt / cuda / Forward 0.000002336 s 0.000002304 s 1.01
cache / PartOpt / cuda / Forward 0.0000023670000000000004 s 0.000002304 s 1.03
cache / IPartOpt / cuda / Forward 0.000002304 s 0.000002273 s 1.01
cache / DefOpt / cuda / Forward 0.000002304 s 0.00000224 s 1.03
cache / IDefOpt / cuda / Forward 0.000002272 s 0.00000224 s 1.01
cache / JaXPipe / cuda / PreRev 0.000011648 s 0.000012128 s 0.96
cache / JaXPipe / cuda / PostRev 0.000011809 s 0.000011712 s 1.01
cache / JaXPipe / cuda / BothRev 0.000011713 s 0.000012096 s 0.97
cache / Jax / cuda / BothRev 0.000011520000000000002 s 0.000012224 s 0.94
cache / HLOOpt / cuda / PreRev 0.000013216 s 0.000013408 s 0.99
cache / HLOOpt / cuda / PostRev 0.000013184 s 0.000013408 s 0.98
cache / HLOOpt / cuda / BothRev 0.000013216 s 0.000013376 s 0.99
cache / PartOpt / cuda / PreRev 0.000011584 s 0.00001216 s 0.95
cache / PartOpt / cuda / PostRev 0.00001184 s 0.000011872 s 1.00
cache / PartOpt / cuda / BothRev 0.000012224 s 0.000012192 s 1.00
cache / IPartOpt / cuda / PreRev 0.000011616 s 0.000012032 s 0.97
cache / IPartOpt / cuda / PostRev 0.000012192 s 0.000012128 s 1.01
cache / IPartOpt / cuda / BothRev 0.000013664 s 0.000011936 s 1.14
cache / DefOpt / cuda / PreRev 0.000011872 s 0.000011904 s 1.00
cache / DefOpt / cuda / PostRev 0.00001376 s 0.000012415 s 1.11
cache / DefOpt / cuda / BothRev 0.000011648 s 0.000011968 s 0.97
cache / IDefOpt / cuda / PreRev 0.000011776 s 0.00001248 s 0.94
cache / IDefOpt / cuda / PostRev 0.000011936 s 0.00001184 s 1.01
cache / IDefOpt / cuda / BothRev 0.000011808 s 0.000011585 s 1.02
cache / JaXPipe / tpu / Primal 0.000002450725 s 0.000002484525 s 0.99
cache / Jax / tpu / Primal 0.0000024579250000000003 s 0.00000246965 s 1.00
cache / HLOOpt / tpu / Primal 0.000002456625 s 0.0000024694 s 0.99
cache / PartOpt / tpu / Primal 0.000002457525 s 0.00000247735 s 0.99
cache / IPartOpt / tpu / Primal 0.000002460675 s 0.000002480875 s 0.99
cache / DefOpt / tpu / Primal 0.000002476125 s 0.0000024787 s 1.00
cache / IDefOpt / tpu / Primal 0.0000024819250000000004 s 0.000002475 s 1.00
cache / JaXPipe / tpu / Forward 0.000003558325 s 0.000003554925 s 1.00
cache / Jax / tpu / Forward 0.0000035434 s 0.000003544425 s 1.00
cache / HLOOpt / tpu / Forward 0.0000035648750000000003 s 0.0000035554249999999995 s 1.00
cache / PartOpt / tpu / Forward 0.000003532 s 0.000003534275 s 1.00
cache / IPartOpt / tpu / Forward 0.000003556075 s 0.000003559375 s 1.00
cache / DefOpt / tpu / Forward 0.0000035440500000000004 s 0.000003535525 s 1.00
cache / IDefOpt / tpu / Forward 0.0000035554000000000003 s 0.0000035686 s 1.00
cache / JaXPipe / tpu / PreRev 0.000004970025 s 0.00000499705 s 0.99
cache / JaXPipe / tpu / PostRev 0.00000497235 s 0.0000050119 s 0.99
cache / JaXPipe / tpu / BothRev 0.000004971025 s 0.0000050275000000000006 s 0.99
cache / Jax / tpu / BothRev 0.00000498565 s 0.000005022575000000001 s 0.99
cache / HLOOpt / tpu / PreRev 0.0000039358250000000006 s 0.00000396875 s 0.99
cache / HLOOpt / tpu / PostRev 0.0000041236 s 0.000004135925 s 1.00
cache / HLOOpt / tpu / BothRev 0.0000039416 s 0.000003967675 s 0.99
cache / PartOpt / tpu / PreRev 0.0000049834 s 0.0000050339 s 0.99
cache / PartOpt / tpu / PostRev 0.000004983150000000001 s 0.00000499745 s 1.00
cache / PartOpt / tpu / BothRev 0.000004975725 s 0.0000050381 s 0.99
cache / IPartOpt / tpu / PreRev 0.000004959324999999999 s 0.000005014575 s 0.99
cache / IPartOpt / tpu / PostRev 0.000004960125 s 0.0000050309 s 0.99
cache / IPartOpt / tpu / BothRev 0.000004980325 s 0.000005007975 s 0.99
cache / DefOpt / tpu / PreRev 0.000004986375 s 0.00000499845 s 1.00
cache / DefOpt / tpu / PostRev 0.0000049932 s 0.0000050070000000000005 s 1.00
cache / DefOpt / tpu / BothRev 0.000004973925000000001 s 0.0000050146 s 0.99
cache / IDefOpt / tpu / PreRev 0.000004974799999999999 s 0.0000050151 s 0.99
cache / IDefOpt / tpu / PostRev 0.0000049811 s 0.000005013374999999999 s 0.99
cache / IDefOpt / tpu / BothRev 0.000004963725 s 0.000005023075 s 0.99
cache / JaXPipe / cpu / Primal 0.000012631 s 0.000007311580002351548 s 1.73
cache / Jax / cpu / Primal 0.000012678 s 0.000008188780020645936 s 1.55
cache / HLOOpt / cpu / Primal 0.00001251 s 0.000007291099991562077 s 1.72
cache / PartOpt / cpu / Primal 0.000012096 s 0.000007267500013767858 s 1.66
cache / IPartOpt / cpu / Primal 0.000012298 s 0.000007347700038735638 s 1.67
cache / DefOpt / cpu / Primal 0.000012347 s 0.000007079820006765658 s 1.74
cache / IDefOpt / cpu / Primal 0.000012303 s 0.000007294299994100584 s 1.69
cache / JaXPipe / cpu / Forward 0.000016492 s 0.000015604340023855912 s 1.06
cache / Jax / cpu / Forward 0.000016797 s 0.000016084959997897386 s 1.04
cache / HLOOpt / cpu / Forward 0.000016706 s 0.000016539119997105445 s 1.01
cache / PartOpt / cpu / Forward 0.000016328 s 0.00002018942001996038 s 0.81
cache / IPartOpt / cpu / Forward 0.000016527 s 0.00001527753997834225 s 1.08
cache / DefOpt / cpu / Forward 0.00001651 s 0.00002051123997262039 s 0.80
cache / IDefOpt / cpu / Forward 0.000016537 s 0.00001486773998294666 s 1.11
cache / JaXPipe / cpu / PreRev 0.000017141 s 0.000016669340011503665 s 1.03
cache / JaXPipe / cpu / PostRev 0.000020367 s 0.000022038719989723177 s 0.92
cache / JaXPipe / cpu / BothRev 0.000017297 s 0.00001681013998677372 s 1.03
cache / Jax / cpu / BothRev 0.000019793 s 0.000022604220002904186 s 0.88
cache / HLOOpt / cpu / PreRev 0.000017277 s 0.000017696700033411617 s 0.98
cache / HLOOpt / cpu / PostRev 0.000016995 s 0.00002060288000393484 s 0.82
cache / HLOOpt / cpu / BothRev 0.000017057 s 0.0000200398199740448 s 0.85
cache / PartOpt / cpu / PreRev 0.000016609 s 0.000017317400006504615 s 0.96
cache / PartOpt / cpu / PostRev 0.000020167 s 0.00002674846000445541 s 0.75
cache / PartOpt / cpu / BothRev 0.000017079 s 0.00001720827998724417 s 0.99
cache / IPartOpt / cpu / PreRev 0.000017274 s 0.00001730513999973482 s 1.00
cache / IPartOpt / cpu / PostRev 0.000017947000000000003 s 0.0000223068999730458 s 0.80
cache / IPartOpt / cpu / BothRev 0.000017746000000000003 s 0.000017066479967979832 s 1.04
cache / DefOpt / cpu / PreRev 0.000016881 s 0.000017424999950890195 s 0.97
cache / DefOpt / cpu / PostRev 0.000017385 s 0.000017306580002696136 s 1.00
cache / DefOpt / cpu / BothRev 0.000017801 s 0.00001792524001757556 s 0.99
cache / IDefOpt / cpu / PreRev 0.000016974 s 0.00001806160001251556 s 0.94
cache / IDefOpt / cpu / PostRev 0.00001736 s 0.00001825745999667561 s 0.95
cache / IDefOpt / cpu / BothRev 0.000017462 s 0.00001824015998863615 s 0.96
Concat / JaXPipe / cpu / Primal 0.00000829088000500633 s 0.000007725860023128916 s 1.07
Concat / Jax / cpu / Primal 0.000007710180016147206 s 0.000007840020052753972 s 0.98
Concat / HLOOpt / cpu / Primal 0.000011147360046379618 s 0.00000750840002183395 s 1.48
Concat / PartOpt / cpu / Primal 0.000007130560006771702 s 0.000007161900002756738 s 1.00
Concat / IPartOpt / cpu / Primal 0.000007127580020096503 s 0.000007567560014649644 s 0.94
Concat / DefOpt / cpu / Primal 0.00001231446000019787 s 0.000011547119993338128 s 1.07
Concat / IDefOpt / cpu / Primal 0.000007326800005102996 s 0.00000715518004653859 s 1.02
Concat / JaXPipe / cpu / Forward 0.000010803660015881178 s 0.000010496600007172674 s 1.03
Concat / Jax / cpu / Forward 0.000010762960009742528 s 0.000011853659989355949 s 0.91
Concat / HLOOpt / cpu / Forward 0.00001643505999709305 s 0.000015137940035856443 s 1.09
Concat / PartOpt / cpu / Forward 0.000016289060013150448 s 0.000015778240040162928 s 1.03
Concat / IPartOpt / cpu / Forward 0.000010891240017372184 s 0.00001113319998694351 s 0.98
Concat / DefOpt / cpu / Forward 0.000015586340014124288 s 0.00001572293999743124 s 0.99
Concat / IDefOpt / cpu / Forward 0.000010860020001928204 s 0.00001081772003999504 s 1.00
Concat / JaXPipe / cpu / PreRev 0.000012803880008505076 s 0.000012534260013126189 s 1.02
Concat / JaXPipe / cpu / PostRev 0.000012497760026235483 s 0.000013171179989512894 s 0.95
Concat / JaXPipe / cpu / BothRev 0.000012766659956469084 s 0.000015525199987678207 s 0.82
Concat / Jax / cpu / BothRev 0.00001265713996872364 s 0.00001536104002298089 s 0.82
Concat / HLOOpt / cpu / PreRev 0.000012894859992229613 s 0.000012696619996859228 s 1.02
Concat / HLOOpt / cpu / PostRev 0.00001647298001444142 s 0.000016494120009156175 s 1.00
Concat / HLOOpt / cpu / BothRev 0.0000144889000148396 s 0.000014265039962992887 s 1.02
Concat / PartOpt / cpu / PreRev 0.00001224205996550154 s 0.000012370900021778652 s 0.99
Concat / PartOpt / cpu / PostRev 0.000012685439978668 s 0.000012754700019286247 s 0.99
Concat / PartOpt / cpu / BothRev 0.000012370480035315268 s 0.00001251313999091508 s 0.99
Concat / IPartOpt / cpu / PreRev 0.000016004760036594236 s 0.000015078859996719985 s 1.06
Concat / IPartOpt / cpu / PostRev 0.000012915980005345772 s 0.000012710540013358696 s 1.02
Concat / IPartOpt / cpu / BothRev 0.000012299160043767188 s 0.00001284707999730017 s 0.96
Concat / DefOpt / cpu / PreRev 0.00001227126002049772 s 0.000012832560005335835 s 0.96
Concat / DefOpt / cpu / PostRev 0.000012404379958752544 s 0.000012438900012057274 s 1.00
Concat / DefOpt / cpu / BothRev 0.0000131081999734306 s 0.000012456539980121309 s 1.05
Concat / IDefOpt / cpu / PreRev 0.000012776739968103357 s 0.000012725619917546284 s 1.00
Concat / IDefOpt / cpu / PostRev 0.000012968260007255594 s 0.000012815459976991406 s 1.01
Concat / IDefOpt / cpu / BothRev 0.000012389780040393815 s 0.000016934820041569765 s 0.73
Concat / JaXPipe / cuda / Primal 0.000001951 s 0.0000019200000000000003 s 1.02
Concat / Jax / cuda / Primal 0.000001951 s 0.0000019200000000000003 s 1.02
Concat / HLOOpt / cuda / Primal 0.000001951 s 0.0000019200000000000003 s 1.02
Concat / PartOpt / cuda / Primal 0.000001951 s 0.0000019200000000000003 s 1.02
Concat / IPartOpt / cuda / Primal 0.000001951 s 0.0000019200000000000003 s 1.02
Concat / DefOpt / cuda / Primal 0.000001951 s 0.0000019200000000000003 s 1.02
Concat / IDefOpt / cuda / Primal 0.000001951 s 0.0000019200000000000003 s 1.02
Concat / JaXPipe / cuda / Forward 0.000010304 s 0.000010752 s 0.96
Concat / Jax / cuda / Forward 0.000010208 s 0.000010945 s 0.93
Concat / HLOOpt / cuda / Forward 0.00001072 s 0.000010592 s 1.01
Concat / PartOpt / cuda / Forward 0.00001008 s 0.000010176 s 0.99
Concat / IPartOpt / cuda / Forward 0.00001008 s 0.000010528 s 0.96
Concat / DefOpt / cuda / Forward 0.000010145 s 0.00001072 s 0.95
Concat / IDefOpt / cuda / Forward 0.000010112 s 0.000010272 s 0.98
Concat / JaXPipe / cuda / PreRev 0.000016768000000000003 s 0.000017344 s 0.97
Concat / JaXPipe / cuda / PostRev 0.000016608 s 0.000017664 s 0.94
Concat / JaXPipe / cuda / BothRev 0.000017055000000000002 s 0.000017568000000000002 s 0.97
Concat / Jax / cuda / BothRev 0.00001664 s 0.000017632 s 0.94
Concat / HLOOpt / cuda / PreRev 0.000017152 s 0.0000176 s 0.97
Concat / HLOOpt / cuda / PostRev 0.000016609 s 0.000017696 s 0.94
Concat / HLOOpt / cuda / BothRev 0.000016864 s 0.0000176 s 0.96
Concat / PartOpt / cuda / PreRev 0.000017312 s 0.000017696 s 0.98
Concat / PartOpt / cuda / PostRev 0.0000168 s 0.000017247999999999998 s 0.97
Concat / PartOpt / cuda / BothRev 0.000016768000000000003 s 0.000017344 s 0.97
Concat / IPartOpt / cuda / PreRev 0.000017312 s 0.000017408 s 0.99
Concat / IPartOpt / cuda / PostRev 0.000016992 s 0.000017760000000000003 s 0.96
Concat / IPartOpt / cuda / BothRev 0.000017375999999999998 s 0.000017567 s 0.99
Concat / DefOpt / cuda / PreRev 0.000019232 s 0.000017664 s 1.09
Concat / DefOpt / cuda / PostRev 0.000016383999999999998 s 0.000017311 s 0.95
Concat / DefOpt / cuda / BothRev 0.000016736 s 0.000018112 s 0.92
Concat / IDefOpt / cuda / PreRev 0.000016864 s 0.000019809 s 0.85
Concat / IDefOpt / cuda / PostRev 0.00001664 s 0.00001728 s 0.96
Concat / IDefOpt / cuda / BothRev 0.000016768999999999998 s 0.000017185 s 0.98
Concat / JaXPipe / tpu / Primal 0.0000015293 s 0.000001528425 s 1.00
Concat / Jax / tpu / Primal 0.0000015238250000000002 s 0.0000015207 s 1.00
Concat / HLOOpt / tpu / Primal 0.000001536075 s 0.0000015292 s 1.00
Concat / PartOpt / tpu / Primal 0.000001534075 s 0.00000152395 s 1.01
Concat / IPartOpt / tpu / Primal 0.000001537525 s 0.000001532075 s 1.00
Concat / DefOpt / tpu / Primal 0.00000152215 s 0.00000151895 s 1.00
Concat / IDefOpt / tpu / Primal 0.00000153635 s 0.000001536425 s 1.00
Concat / JaXPipe / tpu / Forward 0.000001578825 s 0.0000015777 s 1.00
Concat / Jax / tpu / Forward 0.0000015493 s 0.000001549 s 1.00
Concat / HLOOpt / tpu / Forward 0.000001569525 s 0.000001571425 s 1.00
Concat / PartOpt / tpu / Forward 0.000001550625 s 0.00000155265 s 1.00
Concat / IPartOpt / tpu / Forward 0.0000015752 s 0.000001581975 s 1.00
Concat / DefOpt / tpu / Forward 0.000001556425 s 0.00000155355 s 1.00
Concat / IDefOpt / tpu / Forward 0.0000015729000000000002 s 0.0000015713 s 1.00
Concat / JaXPipe / tpu / PreRev 0.00000200425 s 0.0000020133 s 1.00
Concat / JaXPipe / tpu / PostRev 0.0000020867 s 0.00000206465 s 1.01
Concat / JaXPipe / tpu / BothRev 0.0000020176250000000004 s 0.00000201115 s 1.00
Concat / Jax / tpu / BothRev 0.000002076625 s 0.000002063925 s 1.01
Concat / HLOOpt / tpu / PreRev 0.000002011225 s 0.00000201565 s 1.00
Concat / HLOOpt / tpu / PostRev 0.000002075125 s 0.0000020716 s 1.00
Concat / HLOOpt / tpu / BothRev 0.0000020036 s 0.00000200735 s 1.00
Concat / PartOpt / tpu / PreRev 0.000002071225 s 0.000002058425 s 1.01
Concat / PartOpt / tpu / PostRev 0.000002009875 s 0.000002012625 s 1.00
Concat / PartOpt / tpu / BothRev 0.00000207185 s 0.0000020648 s 1.00
Concat / IPartOpt / tpu / PreRev 0.000002012425 s 0.000002007575 s 1.00
Concat / IPartOpt / tpu / PostRev 0.000002076575 s 0.00000206805 s 1.00
Concat / IPartOpt / tpu / BothRev 0.000002018425 s 0.000002002575 s 1.01
Concat / DefOpt / tpu / PreRev 0.000002072325 s 0.0000020624000000000004 s 1.00
Concat / DefOpt / tpu / PostRev 0.0000020013 s 0.0000020068250000000005 s 1.00
Concat / DefOpt / tpu / BothRev 0.0000020715 s 0.00000206625 s 1.00
Concat / IDefOpt / tpu / PreRev 0.000002006025 s 0.00000201155 s 1.00
Concat / IDefOpt / tpu / PostRev 0.0000020828749999999995 s 0.00000206665 s 1.01
Concat / IDefOpt / tpu / BothRev 0.00000200315 s 0.000002006475 s 1.00
Concat / JaXPipe / cpu / Primal 0.000012869 s 0.000007725860023128916 s 1.67
Concat / Jax / cpu / Primal 0.000012686 s 0.000007840020052753972 s 1.62
Concat / HLOOpt / cpu / Primal 0.000012641 s 0.00000750840002183395 s 1.68
Concat / PartOpt / cpu / Primal 0.000012773 s 0.000007161900002756738 s 1.78
Concat / IPartOpt / cpu / Primal 0.000012608 s 0.000007567560014649644 s 1.67
Concat / DefOpt / cpu / Primal 0.000012693 s 0.000011547119993338128 s 1.10
Concat / IDefOpt / cpu / Primal 0.000012398 s 0.00000715518004653859 s 1.73
Concat / JaXPipe / cpu / Forward 0.000017312 s 0.000010496600007172674 s 1.65
Concat / Jax / cpu / Forward 0.000017229 s 0.000011853659989355949 s 1.45
Concat / HLOOpt / cpu / Forward 0.000016947 s 0.000015137940035856443 s 1.12
Concat / PartOpt / cpu / Forward 0.0000174 s 0.000015778240040162928 s 1.10
Concat / IPartOpt / cpu / Forward 0.000017256000000000002 s 0.00001113319998694351 s 1.55
Concat / DefOpt / cpu / Forward 0.000017026 s 0.00001572293999743124 s 1.08
Concat / IDefOpt / cpu / Forward 0.000017565999999999997 s 0.00001081772003999504 s 1.62
Concat / JaXPipe / cpu / PreRev 0.000020006 s 0.000012534260013126189 s 1.60
Concat / JaXPipe / cpu / PostRev 0.000019349 s 0.000013171179989512894 s 1.47
Concat / JaXPipe / cpu / BothRev 0.000019273 s 0.000015525199987678207 s 1.24
Concat / Jax / cpu / BothRev 0.000019919 s 0.00001536104002298089 s 1.30
Concat / HLOOpt / cpu / PreRev 0.000019706 s 0.000012696619996859228 s 1.55
Concat / HLOOpt / cpu / PostRev 0.00001946 s 0.000016494120009156175 s 1.18
Concat / HLOOpt / cpu / BothRev 0.000019042 s 0.000014265039962992887 s 1.33
Concat / PartOpt / cpu / PreRev 0.00001982 s 0.000012370900021778652 s 1.60
Concat / PartOpt / cpu / PostRev 0.000019811 s 0.000012754700019286247 s 1.55
Concat / PartOpt / cpu / BothRev 0.000019259 s 0.00001251313999091508 s 1.54
Concat / IPartOpt / cpu / PreRev 0.000019588000000000003 s 0.000015078859996719985 s 1.30
Concat / IPartOpt / cpu / PostRev 0.000019327 s 0.000012710540013358696 s 1.52
Concat / IPartOpt / cpu / BothRev 0.000019317 s 0.00001284707999730017 s 1.50
Concat / DefOpt / cpu / PreRev 0.000019665 s 0.000012832560005335835 s 1.53
Concat / DefOpt / cpu / PostRev 0.00001957 s 0.000012438900012057274 s 1.57
Concat / DefOpt / cpu / BothRev 0.000019482 s 0.000012456539980121309 s 1.56
Concat / IDefOpt / cpu / PreRev 0.000019182 s 0.000012725619917546284 s 1.51
Concat / IDefOpt / cpu / PostRev 0.000019311 s 0.000012815459976991406 s 1.51
Concat / IDefOpt / cpu / BothRev 0.00001938 s 0.000016934820041569765 s 1.14
const_scatter / JaXPipe / cpu / Primal 0.000007276500000443775 s 0.000009149859961326 s 0.80
const_scatter / Jax / cpu / Primal 0.000007634719986526762 s 0.000008083280008577275 s 0.94
const_scatter / HLOOpt / cpu / Primal 0.000007009179998931359 s 0.0000071863200082589175 s 0.98
const_scatter / PartOpt / cpu / Primal 0.000008064419980655656 s 0.000007375739960480132 s 1.09
const_scatter / IPartOpt / cpu / Primal 0.000007689299973208108 s 0.000007060400021146052 s 1.09
const_scatter / DefOpt / cpu / Primal 0.000010524000008445 s 0.000011595319983825904 s 0.91
const_scatter / IDefOpt / cpu / Primal 0.000007187120008893544 s 0.000007401519987979555 s 0.97
const_scatter / JaXPipe / cpu / Forward 0.00001066464001269196 s 0.000010921500024778652 s 0.98
const_scatter / Jax / cpu / Forward 0.000012009100000796025 s 0.000011650580008790711 s 1.03
const_scatter / HLOOpt / cpu / Forward 0.000014394600002560765 s 0.000018608339969432565 s 0.77
const_scatter / PartOpt / cpu / Forward 0.000015211160016406212 s 0.000015195739988485002 s 1.00
const_scatter / IPartOpt / cpu / Forward 0.000011092420008935732 s 0.000010543240014158071 s 1.05
const_scatter / DefOpt / cpu / Forward 0.000015184720014076448 s 0.00001736160001200915 s 0.87
const_scatter / IDefOpt / cpu / Forward 0.000010472619969732475 s 0.00001035087999298412 s 1.01
const_scatter / JaXPipe / cpu / PreRev 0.0002979728199898 s 0.0003018126200458 s 0.99
const_scatter / JaXPipe / cpu / PostRev 0.0002990370799761 s 0.0002955402599673 s 1.01
const_scatter / JaXPipe / cpu / BothRev 0.0003020345799359 s 0.0002844799399917 s 1.06
const_scatter / Jax / cpu / BothRev 0.0002832198000032 s 0.0002837072600505 s 1.00
const_scatter / HLOOpt / cpu / PreRev 0.0002825039599883 s 0.0002846959200269 s 0.99
const_scatter / HLOOpt / cpu / PostRev 0.0002856400200016 s 0.0002833803199246 s 1.01
const_scatter / HLOOpt / cpu / BothRev 0.0002878017799685 s 0.0002847788800045 s 1.01
const_scatter / PartOpt / cpu / PreRev 0.0002854485200168 s 0.0002836582799682 s 1.01
const_scatter / PartOpt / cpu / PostRev 0.0002853723799671 s 0.0002841413000351 s 1.00
const_scatter / PartOpt / cpu / BothRev 0.0002823793400148 s 0.0002875587399739 s 0.98
const_scatter / IPartOpt / cpu / PreRev 0.0002838197799974 s 0.0002839712200147 s 1.00
const_scatter / IPartOpt / cpu / PostRev 0.0002880095000091 s 0.0002899488999992 s 0.99
const_scatter / IPartOpt / cpu / BothRev 0.0002821923200463 s 0.0002884466199793 s 0.98
const_scatter / DefOpt / cpu / PreRev 0.0002888534600424 s 0.0002823218399953 s 1.02
const_scatter / DefOpt / cpu / PostRev 0.0002890846999889 s 0.0002883031800047 s 1.00
const_scatter / DefOpt / cpu / BothRev 0.0002825164599926 s 0.0002947299199695 s 0.96
const_scatter / IDefOpt / cpu / PreRev 0.0002859800399983 s 0.000284863320021 s 1.00
const_scatter / IDefOpt / cpu / PostRev 0.0002894197600107 s 0.0002873838200321 s 1.01
const_scatter / IDefOpt / cpu / BothRev 0.0002830884000559 s 0.0002888906400039 s 0.98
const_scatter / JaXPipe / cuda / Primal 0.0000019200000000000003 s 0.000001887 s 1.02
const_scatter / Jax / cuda / Primal 0.000001888 s 0.000001887 s 1.00
const_scatter / HLOOpt / cuda / Primal 0.000001888 s 0.000001887 s 1.00
const_scatter / PartOpt / cuda / Primal 0.0000019200000000000003 s 0.000001887 s 1.02
const_scatter / IPartOpt / cuda / Primal 0.000001888 s 0.000001887 s 1.00
const_scatter / DefOpt / cuda / Primal 0.0000019200000000000003 s 0.000001887 s 1.02
const_scatter / IDefOpt / cuda / Primal 0.0000019200000000000003 s 0.000001887 s 1.02
const_scatter / JaXPipe / cuda / Forward 0.000009824 s 0.000011968 s 0.82
const_scatter / Jax / cuda / Forward 0.000010048 s 0.000011776 s 0.85
const_scatter / HLOOpt / cuda / Forward 0.000010048 s 0.000010176 s 0.99
const_scatter / PartOpt / cuda / Forward 0.000010016 s 0.000011744 s 0.85
const_scatter / IPartOpt / cuda / Forward 0.000009856 s 0.000009824 s 1.00
const_scatter / DefOpt / cuda / Forward 0.000009984 s 0.000011904 s 0.84
const_scatter / IDefOpt / cuda / Forward 0.000010208 s 0.00001024 s 1.00
const_scatter / JaXPipe / cuda / PreRev 0.000013184 s 0.000013152 s 1.00
const_scatter / JaXPipe / cuda / PostRev 0.000016992 s 0.000017856 s 0.95
const_scatter / JaXPipe / cuda / BothRev 0.000012992 s 0.000013408 s 0.97
const_scatter / Jax / cuda / BothRev 0.00001712 s 0.000017792 s 0.96
const_scatter / HLOOpt / cuda / PreRev 0.000012992 s 0.000013152 s 0.99
const_scatter / HLOOpt / cuda / PostRev 0.000013088 s 0.000013504 s 0.97
const_scatter / HLOOpt / cuda / BothRev 0.000014016 s 0.000013312 s 1.05
const_scatter / PartOpt / cuda / PreRev 0.000014495 s 0.000013664 s 1.06
const_scatter / PartOpt / cuda / PostRev 0.000018784 s 0.0000176 s 1.07
const_scatter / PartOpt / cuda / BothRev 0.000012864 s 0.00001328 s 0.97
const_scatter / IPartOpt / cuda / PreRev 0.000012993 s 0.000014816 s 0.88
const_scatter / IPartOpt / cuda / PostRev 0.000016993 s 0.00001824 s 0.93
const_scatter / IPartOpt / cuda / BothRev 0.000012608 s 0.000012992 s 0.97
const_scatter / DefOpt / cuda / PreRev 0.0000128 s 0.000012928 s 0.99
const_scatter / DefOpt / cuda / PostRev 0.000012671 s 0.00001328 s 0.95
const_scatter / DefOpt / cuda / BothRev 0.000013024 s 0.000013216 s 0.99
const_scatter / IDefOpt / cuda / PreRev 0.0000128 s 0.00001312 s 0.98
const_scatter / IDefOpt / cuda / PostRev 0.000012928 s 0.000013568 s 0.95
const_scatter / IDefOpt / cuda / BothRev 0.000013024 s 0.00001296 s 1.00
const_scatter / JaXPipe / tpu / Primal 0.000003792925 s 0.000003803325 s 1.00
const_scatter / Jax / tpu / Primal 0.000003819075 s 0.0000038103 s 1.00
const_scatter / HLOOpt / tpu / Primal 9.427e-7 s 9.63325e-7 s 0.98
const_scatter / PartOpt / tpu / Primal 0.000003830225 s 0.0000038188 s 1.00
const_scatter / IPartOpt / tpu / Primal 0.0000038086 s 0.000003788875 s 1.01
const_scatter / DefOpt / tpu / Primal 9.74075e-7 s 9.70025e-7 s 1.00
const_scatter / IDefOpt / tpu / Primal 9.422e-7 s 9.495e-7 s 0.99
const_scatter / JaXPipe / tpu / Forward 0.0000019348 s 0.000001921075 s 1.01
const_scatter / Jax / tpu / Forward 0.000006473075 s 0.000006476150000000001 s 1.00
const_scatter / HLOOpt / tpu / Forward 0.00000192395 s 0.000001942025 s 0.99
const_scatter / PartOpt / tpu / Forward 0.00000193495 s 0.000001940325 s 1.00
const_scatter / IPartOpt / tpu / Forward 0.000001928225 s 0.00000192275 s 1.00
const_scatter / DefOpt / tpu / Forward 0.000001927725 s 0.000001944975 s 0.99
const_scatter / IDefOpt / tpu / Forward 0.000001912525 s 0.000001944275 s 0.98
const_scatter / JaXPipe / tpu / PreRev 0.0000043012250000000005 s 0.000004326624999999999 s 0.99
const_scatter / JaXPipe / tpu / PostRev 0.0000066198000000000006 s 0.000006591275 s 1.00
const_scatter / JaXPipe / tpu / BothRev 0.00000429755 s 0.000004323449999999999 s 0.99
const_scatter / Jax / tpu / BothRev 0.000006632925 s 0.000006607575 s 1.00
const_scatter / HLOOpt / tpu / PreRev 0.00000428535 s 0.0000043108 s 0.99
const_scatter / HLOOpt / tpu / PostRev 0.000004296675 s 0.0000043236 s 0.99
const_scatter / HLOOpt / tpu / BothRev 0.0000042939250000000006 s 0.000004320025 s 0.99
const_scatter / PartOpt / tpu / PreRev 0.000004294975 s 0.0000043140000000000005 s 1.00
const_scatter / PartOpt / tpu / PostRev 0.000006596375 s 0.0000066094 s 1.00
const_scatter / PartOpt / tpu / BothRev 0.000004289425 s 0.000004309574999999999 s 1.00
const_scatter / IPartOpt / tpu / PreRev 0.000004285625 s 0.000004329225 s 0.99
const_scatter / IPartOpt / tpu / PostRev 0.000006614725 s 0.000006611575 s 1.00
const_scatter / IPartOpt / tpu / BothRev 0.0000042748 s 0.000004309999999999999 s 0.99
const_scatter / DefOpt / tpu / PreRev 0.000004305525 s 0.0000043274 s 0.99
const_scatter / DefOpt / tpu / PostRev 0.000004298424999999999 s 0.0000043121 s 1.00
const_scatter / DefOpt / tpu / BothRev 0.000004301825 s 0.00000433185 s 0.99
const_scatter / IDefOpt / tpu / PreRev 0.000004290725 s 0.0000043074000000000005 s 1.00
const_scatter / IDefOpt / tpu / PostRev 0.00000428945 s 0.0000043165 s 0.99
const_scatter / IDefOpt / tpu / BothRev 0.0000042994 s 0.0000043184 s 1.00
const_scatter / JaXPipe / cpu / Primal 0.000012786 s 0.000009149859961326 s 1.40
const_scatter / Jax / cpu / Primal 0.000012586 s 0.000008083280008577275 s 1.56
const_scatter / HLOOpt / cpu / Primal 0.00001258 s 0.0000071863200082589175 s 1.75
const_scatter / PartOpt / cpu / Primal 0.000012503 s 0.000007375739960480132 s 1.70
const_scatter / IPartOpt / cpu / Primal 0.000012659 s 0.000007060400021146052 s 1.79
const_scatter / DefOpt / cpu / Primal 0.000012324 s 0.000011595319983825904 s 1.06
const_scatter / IDefOpt / cpu / Primal 0.000012411 s 0.000007401519987979555 s 1.68
const_scatter / JaXPipe / cpu / Forward 0.000017007 s 0.000010921500024778652 s 1.56
const_scatter / Jax / cpu / Forward 0.000027415 s 0.000011650580008790711 s 2.35
const_scatter / HLOOpt / cpu / Forward 0.000016445000000000003 s 0.000018608339969432565 s 0.88
const_scatter / PartOpt / cpu / Forward 0.000016205 s 0.000015195739988485002 s 1.07
const_scatter / IPartOpt / cpu / Forward 0.000016742 s 0.000010543240014158071 s 1.59
const_scatter / DefOpt / cpu / Forward 0.00001658 s 0.00001736160001200915 s 0.95
const_scatter / IDefOpt / cpu / Forward 0.000016499 s 0.00001035087999298412 s 1.59
const_scatter / JaXPipe / cpu / PreRev 0.000489043 s 0.0003018126200458 s 1.62
const_scatter / JaXPipe / cpu / PostRev 0.000485883 s 0.0002955402599673 s 1.64
const_scatter / JaXPipe / cpu / BothRev 0.000500807 s 0.0002844799399917 s 1.76
const_scatter / Jax / cpu / BothRev 0.000508842 s 0.0002837072600505 s 1.79
const_scatter / HLOOpt / cpu / PreRev 0.000506338 s 0.0002846959200269 s 1.78
const_scatter / HLOOpt / cpu / PostRev 0.000513722 s 0.0002833803199246 s 1.81
const_scatter / HLOOpt / cpu / BothRev 0.000507349 s 0.0002847788800045 s 1.78
const_scatter / PartOpt / cpu / PreRev 0.000491381 s 0.0002836582799682 s 1.73
const_scatter / PartOpt / cpu / PostRev 0.000498466 s 0.0002841413000351 s 1.75
const_scatter / PartOpt / cpu / BothRev 0.000515623 s 0.0002875587399739 s 1.79
const_scatter / IPartOpt / cpu / PreRev 0.000512499 s 0.0002839712200147 s 1.80
const_scatter / IPartOpt / cpu / PostRev 0.000494 s 0.0002899488999992 s 1.70
const_scatter / IPartOpt / cpu / BothRev 0.000503447 s 0.0002884466199793 s 1.75
const_scatter / DefOpt / cpu / PreRev 0.000504931 s 0.0002823218399953 s 1.79
const_scatter / DefOpt / cpu / PostRev 0.000507416 s 0.0002883031800047 s 1.76
const_scatter / DefOpt / cpu / BothRev 0.000495143 s 0.0002947299199695 s 1.68
const_scatter / IDefOpt / cpu / PreRev 0.000493116 s 0.000284863320021 s 1.73
const_scatter / IDefOpt / cpu / PostRev 0.000507795 s 0.0002873838200321 s 1.77
const_scatter / IDefOpt / cpu / BothRev 0.000495178 s 0.0002888906400039 s 1.71
GenDot / JaXPipe / cpu / Primal 0.00000915183996767155 s 0.000010176719997616602 s 0.90
GenDot / Jax / cpu / Primal 0.000007920060015749186 s 0.000008191240021915292 s 0.97
GenDot / HLOOpt / cpu / Primal 0.00001297285995860875 s 0.000013252619964987389 s 0.98
GenDot / PartOpt / cpu / Primal 0.000008386560029975953 s 0.000008220199997595046 s 1.02
GenDot / IPartOpt / cpu / Primal 0.00000832554001135577 s 0.00000903473997823312 s 0.92
GenDot / DefOpt / cpu / Primal 0.00000840552000227035 s 0.000012143100011599016 s 0.69
GenDot / IDefOpt / cpu / Primal 0.000008828000018183957 s 0.000008335859993167105 s 1.06
GenDot / JaXPipe / cpu / Forward 0.000012322299980951356 s 0.000012309759986237625 s 1.00
GenDot / Jax / cpu / Forward 0.000010743180000645223 s 0.000011215760005143236 s 0.96
GenDot / HLOOpt / cpu / Forward 0.000012437360010153495 s 0.00001224884002112958 s 1.02
GenDot / PartOpt / cpu / Forward 0.00001738322004712245 s 0.000012817179995181504 s 1.36
GenDot / IPartOpt / cpu / Forward 0.000011878219984282623 s 0.0000129771400224854 s 0.92
GenDot / DefOpt / cpu / Forward 0.00001744481999594427 s 0.00001729242000692466 s 1.01
GenDot / IDefOpt / cpu / Forward 0.0000120781999976316 s 0.000011872980030602775 s 1.02
GenDot / JaXPipe / cpu / PreRev 0.00001233391998539446 s 0.000011857600020448444 s 1.04
GenDot / JaXPipe / cpu / PostRev 0.000010956680043818778 s 0.000011128899986942997 s 0.98
GenDot / JaXPipe / cpu / BothRev 0.000013263999926493853 s 0.000012159080024503054 s 1.09
GenDot / Jax / cpu / BothRev 0.000011447280039647012 s 0.000011684280007102644 s 0.98
GenDot / HLOOpt / cpu / PreRev 0.00001198541998746805 s 0.0000117795399910392 s 1.02
GenDot / HLOOpt / cpu / PostRev 0.000016266200018435483 s 0.000011763639986384076 s 1.38
GenDot / HLOOpt / cpu / BothRev 0.00001332410002760298 s 0.000013835919953635313 s 0.96
GenDot / PartOpt / cpu / PreRev 0.00001254238000001351 s 0.000012417100015227332 s 1.01
GenDot / PartOpt / cpu / PostRev 0.000011660099999062369 s 0.000011310380014037946 s 1.03
GenDot / PartOpt / cpu / BothRev 0.000012198319991512108 s 0.000011758400014514336 s 1.04
GenDot / IPartOpt / cpu / PreRev 0.00001774587994077592 s 0.000013740780032094337 s 1.29
GenDot / IPartOpt / cpu / PostRev 0.000011386940004740608 s 0.000012285620023249068 s 0.93
GenDot / IPartOpt / cpu / BothRev 0.000011752620030165418 s 0.000011663839995890157 s 1.01
GenDot / DefOpt / cpu / PreRev 0.000011530220026543248 s 0.000012125860048399772 s 0.95
GenDot / DefOpt / cpu / PostRev 0.000011647239989542869 s 0.000011902560027010624 s 0.98
GenDot / DefOpt / cpu / BothRev 0.000011774300000979563 s 0.000012017819954053264 s 0.98
GenDot / IDefOpt / cpu / PreRev 0.000011579380025068531 s 0.000012305520049267216 s 0.94
GenDot / IDefOpt / cpu / PostRev 0.000012202139987493866 s 0.00001163964001534623 s 1.05
GenDot / IDefOpt / cpu / BothRev 0.000011632200021267635 s 0.000012038459981340566 s 0.97
GenDot / JaXPipe / cuda / Primal 0.000002016 s 0.000002016 s 1
GenDot / Jax / cuda / Primal 0.000002016 s 0.000002015 s 1.00
GenDot / HLOOpt / cuda / Primal 0.000002015 s 0.000001984 s 1.02
GenDot / PartOpt / cuda / Primal 0.000002016 s 0.000002016 s 1
GenDot / IPartOpt / cuda / Primal 0.000002016 s 0.000002015 s 1.00
GenDot / DefOpt / cuda / Primal 0.000002016 s 0.000001984 s 1.02
GenDot / IDefOpt / cuda / Primal 0.000002016 s 0.000001984 s 1.02
GenDot / JaXPipe / cuda / Forward 0.000009825 s 0.0000104 s 0.94
GenDot / Jax / cuda / Forward 0.00001024 s 0.000010273 s 1.00
GenDot / HLOOpt / cuda / Forward 0.000010432 s 0.00001024 s 1.02
GenDot / PartOpt / cuda / Forward 0.000009984 s 0.000010208 s 0.98
GenDot / IPartOpt / cuda / Forward 0.000010336 s 0.000010432 s 0.99
GenDot / DefOpt / cuda / Forward 0.00001008 s 0.000010528 s 0.96
GenDot / IDefOpt / cuda / Forward 0.00001024 s 0.000010272 s 1.00
GenDot / JaXPipe / cuda / PreRev 0.000010272 s 0.000010464 s 0.98
GenDot / JaXPipe / cuda / PostRev 0.000010271 s 0.00001072 s 0.96
GenDot / JaXPipe / cuda / BothRev 0.000010336 s 0.000010113 s 1.02
GenDot / Jax / cuda / BothRev 0.000010272 s 0.000010592 s 0.97
GenDot / HLOOpt / cuda / PreRev 0.000010016 s 0.000010304 s 0.97
GenDot / HLOOpt / cuda / PostRev 0.000010048 s 0.000010592 s 0.95
GenDot / HLOOpt / cuda / BothRev 0.000010176 s 0.000009984 s 1.02
GenDot / PartOpt / cuda / PreRev 0.000010111 s 0.000010464 s 0.97
GenDot / PartOpt / cuda / PostRev 0.000010528 s 0.000010336 s 1.02
GenDot / PartOpt / cuda / BothRev 0.000010272 s 0.000010432 s 0.98
GenDot / IPartOpt / cuda / PreRev 0.000010144 s 0.000010656 s 0.95
GenDot / IPartOpt / cuda / PostRev 0.000010304 s 0.000010336 s 1.00
GenDot / IPartOpt / cuda / BothRev 0.000010144 s 0.000009984 s 1.02
GenDot / DefOpt / cuda / PreRev 0.000010016 s 0.000010433 s 0.96
GenDot / DefOpt / cuda / PostRev 0.000009824 s 0.000010624 s 0.92
GenDot / DefOpt / cuda / BothRev 0.000010207 s 0.00001056 s 0.97
GenDot / IDefOpt / cuda / PreRev 0.000010112 s 0.000010304 s 0.98
GenDot / IDefOpt / cuda / PostRev 0.000010528 s 0.000010944 s 0.96
GenDot / IDefOpt / cuda / BothRev 0.000010047 s 0.000010528 s 0.95
GenDot / JaXPipe / tpu / Primal 9.3015e-7 s 9.2645e-7 s 1.00
GenDot / Jax / tpu / Primal 9.351e-7 s 9.35875e-7 s 1.00
GenDot / HLOOpt / tpu / Primal 0.0000015831 s 0.00000155025 s 1.02
GenDot / PartOpt / tpu / Primal 9.361e-7 s 9.3565e-7 s 1.00
GenDot / IPartOpt / tpu / Primal 9.39675e-7 s 9.3585e-7 s 1.00
GenDot / DefOpt / tpu / Primal 0.0000014862 s 0.00000148595 s 1.00
GenDot / IDefOpt / tpu / Primal 0.0000015700249999999998 s 0.0000015580250000000002 s 1.01
GenDot / JaXPipe / tpu / Forward 0.0000031503 s 0.000003161325 s 1.00
GenDot / Jax / tpu / Forward 0.000002326075 s 0.000002324375 s 1.00
GenDot / HLOOpt / tpu / Forward 0.00000311805 s 0.0000031084 s 1.00
GenDot / PartOpt / tpu / Forward 0.00000321715 s 0.000003208625 s 1.00
GenDot / IPartOpt / tpu / Forward 0.000003106175 s 0.0000031077 s 1.00
GenDot / DefOpt / tpu / Forward 0.0000032195249999999995 s 0.0000032057 s 1.00
GenDot / IDefOpt / tpu / Forward 0.0000031099 s 0.000003114375 s 1.00
GenDot / JaXPipe / tpu / PreRev 0.0000029548 s 0.000002948925 s 1.00
GenDot / JaXPipe / tpu / PostRev 0.000002402625 s 0.0000024131500000000003 s 1.00
GenDot / JaXPipe / tpu / BothRev 0.0000029557250000000004 s 0.000002966675 s 1.00
GenDot / Jax / tpu / BothRev 0.0000024020250000000005 s 0.00000240555 s 1.00
GenDot / HLOOpt / tpu / PreRev 0.0000029540250000000005 s 0.00000294995 s 1.00
GenDot / HLOOpt / tpu / PostRev 0.000002923575 s 0.000002928175 s 1.00
GenDot / HLOOpt / tpu / BothRev 0.00000294605 s 0.0000029566 s 1.00
GenDot / PartOpt / tpu / PreRev 0.00000292315 s 0.0000029255 s 1.00
GenDot / PartOpt / tpu / PostRev 0.000002415125 s 0.000002390775 s 1.01
GenDot / PartOpt / tpu / BothRev 0.000002925575 s 0.000002923425 s 1.00
GenDot / IPartOpt / tpu / PreRev 0.00000295385 s 0.000002956225 s 1.00
GenDot / IPartOpt / tpu / PostRev 0.000002417475 s 0.00000241075 s 1.00
GenDot / IPartOpt / tpu / BothRev 0.0000029557 s 0.0000029465000000000005 s 1.00
GenDot / DefOpt / tpu / PreRev 0.0000029258250000000005 s 0.000002928475 s 1.00
GenDot / DefOpt / tpu / PostRev 0.000002958625 s 0.00000294945 s 1.00
GenDot / DefOpt / tpu / BothRev 0.0000029224499999999994 s 0.0000029398750000000004 s 0.99
GenDot / IDefOpt / tpu / PreRev 0.0000029460750000000004 s 0.000002941075 s 1.00
GenDot / IDefOpt / tpu / PostRev 0.00000293185 s 0.0000029294 s 1.00
GenDot / IDefOpt / tpu / BothRev 0.000002957475 s 0.00000294325 s 1.00
GenDot / JaXPipe / cpu / Primal 0.000014514 s 0.000010176719997616602 s 1.43
GenDot / Jax / cpu / Primal 0.000014676 s 0.000008191240021915292 s 1.79
GenDot / HLOOpt / cpu / Primal 0.000013592 s 0.000013252619964987389 s 1.03
GenDot / PartOpt / cpu / Primal 0.000014611 s 0.000008220199997595046 s 1.78
GenDot / IPartOpt / cpu / Primal 0.000022746 s 0.00000903473997823312 s 2.52
GenDot / DefOpt / cpu / Primal 0.000013551 s 0.000012143100011599016 s 1.12
GenDot / IDefOpt / cpu / Primal 0.000013507999999999998 s 0.000008335859993167105 s 1.62
GenDot / JaXPipe / cpu / Forward 0.000019026 s 0.000012309759986237625 s 1.55
GenDot / Jax / cpu / Forward 0.000019446 s 0.000011215760005143236 s 1.73
GenDot / HLOOpt / cpu / Forward 0.000018807 s 0.00001224884002112958 s 1.54
GenDot / PartOpt / cpu / Forward 0.000018697 s 0.000012817179995181504 s 1.46
GenDot / IPartOpt / cpu / Forward 0.000018382 s 0.0000129771400224854 s 1.42
GenDot / DefOpt / cpu / Forward 0.00001891 s 0.00001729242000692466 s 1.09
GenDot / IDefOpt / cpu / Forward 0.000018809 s 0.000011872980030602775 s 1.58
GenDot / JaXPipe / cpu / PreRev 0.000019334 s 0.000011857600020448444 s 1.63
GenDot / JaXPipe / cpu / PostRev 0.000020364 s 0.000011128899986942997 s 1.83
GenDot / JaXPipe / cpu / BothRev 0.000019489 s 0.000012159080024503054 s 1.60
GenDot / Jax / cpu / BothRev 0.000019821 s 0.000011684280007102644 s 1.70
GenDot / HLOOpt / cpu / PreRev 0.000019254 s 0.0000117795399910392 s 1.63
GenDot / HLOOpt / cpu / PostRev 0.00001906 s 0.000011763639986384076 s 1.62
GenDot / HLOOpt / cpu / BothRev 0.000019288 s 0.000013835919953635313 s 1.39
GenDot / PartOpt / cpu / PreRev 0.000018834 s 0.000012417100015227332 s 1.52
GenDot / PartOpt / cpu / PostRev 0.000020615 s 0.000011310380014037946 s 1.82
GenDot / PartOpt / cpu / BothRev 0.000019263 s 0.000011758400014514336 s 1.64
GenDot / IPartOpt / cpu / PreRev 0.000018673 s 0.000013740780032094337 s 1.36
GenDot / IPartOpt / cpu / PostRev 0.000020835 s 0.000012285620023249068 s 1.70
GenDot / IPartOpt / cpu / BothRev 0.000019565 s 0.000011663839995890157 s 1.68
GenDot / DefOpt / cpu / PreRev 0.00001908 s 0.000012125860048399772 s 1.57
GenDot / DefOpt / cpu / PostRev 0.000019228 s 0.000011902560027010624 s 1.62
GenDot / DefOpt / cpu / BothRev 0.000019516 s 0.000012017819954053264 s 1.62
GenDot / IDefOpt / cpu / PreRev 0.000018701 s 0.000012305520049267216 s 1.52
GenDot / IDefOpt / cpu / PostRev 0.000018862 s 0.00001163964001534623 s 1.62
GenDot / IDefOpt / cpu / BothRev 0.000019534 s 0.000012038459981340566 s 1.62
hlo_ffi / JaXPipe / cpu / Primal 0.00001210560001709382 s 0.000011426740011302172 s 1.06
hlo_ffi / Jax / cpu / Primal 0.00001176938004391559 s 0.000011207280040252954 s 1.05
hlo_ffi / HLOOpt / cpu / Primal 0.000013649639977302283 s 0.000014693479970446788 s 0.93
hlo_ffi / PartOpt / cpu / Primal 0.000011161680013174192 s 0.00001102379998883407 s 1.01
hlo_ffi / IPartOpt / cpu / Primal 0.000011487680039863334 s 0.000010534919974816148 s 1.09
hlo_ffi / DefOpt / cpu / Primal 0.00001337863996013766 s 0.000010974839942718972 s 1.22
hlo_ffi / IDefOpt / cpu / Primal 0.000011655660000542413 s 0.00001072148001185269 s 1.09
hlo_ffi / JaXPipe / cpu / Forward 0.000017495640013294177 s 0.00001635991998227837 s 1.07
hlo_ffi / Jax / cpu / Forward 0.000016945179968388402 s 0.00001635847996112716 s 1.04
hlo_ffi / HLOOpt / cpu / Forward 0.00001739833997817186 s 0.00001682942001025367 s 1.03
hlo_ffi / PartOpt / cpu / Forward 0.000017069340019588706 s 0.00001639388000512554 s 1.04
hlo_ffi / IPartOpt / cpu / Forward 0.000016617760002191063 s 0.000016555979991608184 s 1.00
hlo_ffi / DefOpt / cpu / Forward 0.000017595540039110348 s 0.000017256500013900224 s 1.02
hlo_ffi / IDefOpt / cpu / Forward 0.00001729596000586753 s 0.000016987700000754557 s 1.02
hlo_ffi / JaXPipe / cpu / PreRev 0.000016330299977198593 s 0.000015459200012628573 s 1.06
hlo_ffi / JaXPipe / cpu / PostRev 0.00001699681997706648 s 0.0000153673000386334 s 1.11
hlo_ffi / JaXPipe / cpu / BothRev 0.000016836080003486132 s 0.000015725039938843112 s 1.07
hlo_ffi / Jax / cpu / BothRev 0.000016905639995457022 s 0.000015293040023607317 s 1.11
hlo_ffi / HLOOpt / cpu / PreRev 0.000016722579994166152 s 0.000015225179995468352 s 1.10
hlo_ffi / HLOOpt / cpu / PostRev 0.000016344340010618908 s 0.00001628216003155103 s 1.00
hlo_ffi / HLOOpt / cpu / BothRev 0.000018543020014476498 s 0.000017381139996359708 s 1.07
hlo_ffi / PartOpt / cpu / PreRev 0.000016903819969229516 s 0.000015519260014116298 s 1.09
hlo_ffi / PartOpt / cpu / PostRev 0.00001703957997051475 s 0.000015806079964022502 s 1.08
hlo_ffi / PartOpt / cpu / BothRev 0.000017722480006341357 s 0.000015853700006118743 s 1.12
hlo_ffi / IPartOpt / cpu / PreRev 0.000017461859997638384 s 0.000015407880000566364 s 1.13
hlo_ffi / IPartOpt / cpu / PostRev 0.000017301559973930126 s 0.000015501520019824967 s 1.12
hlo_ffi / IPartOpt / cpu / BothRev 0.000017021039984683738 s 0.00001576622002176009 s 1.08
hlo_ffi / DefOpt / cpu / PreRev 0.00001668742000219936 s 0.000015992219960025978 s 1.04
hlo_ffi / DefOpt / cpu / PostRev 0.000016636140035188873 s 0.000015379819997178858 s 1.08
hlo_ffi / DefOpt / cpu / BothRev 0.000017888880011014406 s 0.000015584500006298184 s 1.15
hlo_ffi / IDefOpt / cpu / PreRev 0.00001683047996266396 s 0.000015718799995738665 s 1.07
hlo_ffi / IDefOpt / cpu / PostRev 0.000016570299994782544 s 0.000015333300034399144 s 1.08
hlo_ffi / IDefOpt / cpu / BothRev 0.00001750478003486933 s 0.000015546340018772754 s 1.13
hlo_ffi / JaXPipe / cuda / Primal 0.000001984 s 0.000001983 s 1.00
hlo_ffi / Jax / cuda / Primal 0.000001984 s 0.000001983 s 1.00
hlo_ffi / HLOOpt / cuda / Primal 0.000001984 s 0.000001983 s 1.00
hlo_ffi / PartOpt / cuda / Primal 0.000001983 s 0.000001983 s 1
hlo_ffi / IPartOpt / cuda / Primal 0.000001984 s 0.000001983 s 1.00
hlo_ffi / DefOpt / cuda / Primal 0.000001984 s 0.000001983 s 1.00
hlo_ffi / IDefOpt / cuda / Primal 0.000001984 s 0.000001984 s 1
hlo_ffi / JaXPipe / cuda / Forward 0.00000208 s 0.00000208 s 1
hlo_ffi / Jax / cuda / Forward 0.00000208 s 0.000002047 s 1.02
hlo_ffi / HLOOpt / cuda / Forward 0.00000208 s 0.00000208 s 1
hlo_ffi / PartOpt / cuda / Forward 0.00000208 s 0.00000208 s 1
hlo_ffi / IPartOpt / cuda / Forward 0.00000208 s 0.000002048 s 1.02
hlo_ffi / DefOpt / cuda / Forward 0.00000208 s 0.000002048 s 1.02
hlo_ffi / IDefOpt / cuda / Forward 0.00000208 s 0.00000208 s 1
hlo_ffi / JaXPipe / cuda / PreRev 0.000002048 s 0.000002048 s 1
hlo_ffi / JaXPipe / cuda / PostRev 0.000002048 s 0.000002047 s 1.00
hlo_ffi / JaXPipe / cuda / BothRev 0.000002048 s 0.000002047 s 1.00
hlo_ffi / Jax / cuda / BothRev 0.000002048 s 0.000002047 s 1.00
hlo_ffi / HLOOpt / cuda / PreRev 0.000002048 s 0.000002047 s 1.00
hlo_ffi / HLOOpt / cuda / PostRev 0.000002047 s 0.000002048 s 1.00
hlo_ffi / HLOOpt / cuda / BothRev 0.00000208 s 0.000002048 s 1.02
hlo_ffi / PartOpt / cuda / PreRev 0.000002048 s 0.000002047 s 1.00
hlo_ffi / PartOpt / cuda / PostRev 0.000002048 s 0.000002047 s 1.00
hlo_ffi / PartOpt / cuda / BothRev 0.000002048 s 0.000002048 s 1
hlo_ffi / IPartOpt / cuda / PreRev 0.000002048 s 0.000002047 s 1.00
hlo_ffi / IPartOpt / cuda / PostRev 0.000002047 s 0.000002048 s 1.00
hlo_ffi / IPartOpt / cuda / BothRev 0.000002048 s 0.000002048 s 1
hlo_ffi / DefOpt / cuda / PreRev 0.000002048 s 0.000002047 s 1.00
hlo_ffi / DefOpt / cuda / PostRev 0.000002048 s 0.000002048 s 1
hlo_ffi / DefOpt / cuda / BothRev 0.00000208 s 0.000002048 s 1.02
hlo_ffi / IDefOpt / cuda / PreRev 0.000002049 s 0.000002048 s 1.00
hlo_ffi / IDefOpt / cuda / PostRev 0.00000208 s 0.000002047 s 1.02
hlo_ffi / IDefOpt / cuda / BothRev 0.00000208 s 0.000002047 s 1.02
hlo_ffi / JaXPipe / tpu / Primal 9.19925e-7 s 9.1895e-7 s 1.00
hlo_ffi / Jax / tpu / Primal 9.555750000000002e-7 s 9.5215e-7 s 1.00
hlo_ffi / HLOOpt / tpu / Primal 8.955e-7 s 9.004e-7 s 0.99
hlo_ffi / PartOpt / tpu / Primal 9.52875e-7 s 9.5305e-7 s 1.00
hlo_ffi / IPartOpt / tpu / Primal 8.988000000000001e-7 s 8.98275e-7 s 1.00
hlo_ffi / DefOpt / tpu / Primal 9.51325e-7 s 9.4945e-7 s 1.00
hlo_ffi / IDefOpt / tpu / Primal 8.9685e-7 s 8.96675e-7 s 1.00
hlo_ffi / JaXPipe / tpu / Forward 9.49e-7 s 9.4955e-7 s 1.00
hlo_ffi / Jax / tpu / Forward 9.82025e-7 s 9.8205e-7 s 1.00
hlo_ffi / HLOOpt / tpu / Forward 9.74e-7 s 9.73925e-7 s 1.00
hlo_ffi / PartOpt / tpu / Forward 9.346e-7 s 9.3375e-7 s 1.00
hlo_ffi / IPartOpt / tpu / Forward 9.743e-7 s 9.74375e-7 s 1.00
hlo_ffi / DefOpt / tpu / Forward 9.34075e-7 s 9.33775e-7 s 1.00
hlo_ffi / IDefOpt / tpu / Forward 9.741e-7 s 9.73975e-7 s 1.00
hlo_ffi / JaXPipe / tpu / PreRev 9.3245e-7 s 9.323e-7 s 1.00
hlo_ffi / JaXPipe / tpu / PostRev 9.646e-7 s 9.6545e-7 s 1.00
hlo_ffi / JaXPipe / tpu / BothRev 9.59925e-7 s 9.60375e-7 s 1.00
hlo_ffi / Jax / tpu / BothRev 9.65225e-7 s 9.65125e-7 s 1.00
hlo_ffi / HLOOpt / tpu / PreRev 9.59925e-7 s 9.598e-7 s 1.00
hlo_ffi / HLOOpt / tpu / PostRev 9.6515e-7 s 9.6545e-7 s 1.00
hlo_ffi / HLOOpt / tpu / BothRev 9.5985e-7 s 9.6045e-7 s 1.00
hlo_ffi / PartOpt / tpu / PreRev 9.649e-7 s 9.65375e-7 s 1.00
hlo_ffi / PartOpt / tpu / PostRev 9.59925e-7 s 9.6025e-7 s 1.00
hlo_ffi / PartOpt / tpu / BothRev 9.64775e-7 s 9.654e-7 s 1.00
hlo_ffi / IPartOpt / tpu / PreRev 9.59525e-7 s 9.59875e-7 s 1.00
hlo_ffi / IPartOpt / tpu / PostRev 9.64875e-7 s 9.65175e-7 s 1.00
hlo_ffi / IPartOpt / tpu / BothRev 9.59775e-7 s 9.6015e-7 s 1.00
hlo_ffi / DefOpt / tpu / PreRev 9.65125e-7 s 9.65225e-7 s 1.00
hlo_ffi / DefOpt / tpu / PostRev 9.59425e-7 s 9.60275e-7 s 1.00
hlo_ffi / DefOpt / tpu / BothRev 9.64875e-7 s 9.652e-7 s 1.00
hlo_ffi / IDefOpt / tpu / PreRev 9.6e-7 s 9.60075e-7 s 1.00
hlo_ffi / IDefOpt / tpu / PostRev 9.6485e-7 s 9.654e-7 s 1.00
hlo_ffi / IDefOpt / tpu / BothRev 9.599e-7 s 9.60125e-7 s 1.00
hlo_ffi / JaXPipe / cpu / Primal 0.00001747 s 0.000011426740011302172 s 1.53
hlo_ffi / Jax / cpu / Primal 0.000017531000000000002 s 0.000011207280040252954 s 1.56
hlo_ffi / HLOOpt / cpu / Primal 0.000017542999999999998 s 0.000014693479970446788 s 1.19
hlo_ffi / PartOpt / cpu / Primal 0.00001757 s 0.00001102379998883407 s 1.59
hlo_ffi / IPartOpt / cpu / Primal 0.000017236000000000002 s 0.000010534919974816148 s 1.64
hlo_ffi / DefOpt / cpu / Primal 0.000017319 s 0.000010974839942718972 s 1.58
hlo_ffi / IDefOpt / cpu / Primal 0.000017303 s 0.00001072148001185269 s 1.61
hlo_ffi / JaXPipe / cpu / Forward 0.000024548 s 0.00001635991998227837 s 1.50
hlo_ffi / Jax / cpu / Forward 0.000024127 s 0.00001635847996112716 s 1.47
hlo_ffi / HLOOpt / cpu / Forward 0.000024032 s 0.00001682942001025367 s 1.43
hlo_ffi / PartOpt / cpu / Forward 0.00002418 s 0.00001639388000512554 s 1.47
hlo_ffi / IPartOpt / cpu / Forward 0.000024529 s 0.000016555979991608184 s 1.48
hlo_ffi / DefOpt / cpu / Forward 0.000024361 s 0.000017256500013900224 s 1.41
hlo_ffi / IDefOpt / cpu / Forward 0.00002442 s 0.000016987700000754557 s 1.44
hlo_ffi / JaXPipe / cpu / PreRev 0.000024409 s 0.000015459200012628573 s 1.58
hlo_ffi / JaXPipe / cpu / PostRev 0.000024939 s 0.0000153673000386334 s 1.62
hlo_ffi / JaXPipe / cpu / BothRev 0.000024563 s 0.000015725039938843112 s 1.56
hlo_ffi / Jax / cpu / BothRev 0.00002383 s 0.000015293040023607317 s 1.56
hlo_ffi / HLOOpt / cpu / PreRev 0.00002427 s 0.000015225179995468352 s 1.59
hlo_ffi / HLOOpt / cpu / PostRev 0.00002512 s 0.00001628216003155103 s 1.54
hlo_ffi / HLOOpt / cpu / BothRev 0.000024284 s 0.000017381139996359708 s 1.40
hlo_ffi / PartOpt / cpu / PreRev 0.000023992 s 0.000015519260014116298 s 1.55
hlo_ffi / PartOpt / cpu / PostRev 0.000025022 s 0.000015806079964022502 s 1.58
hlo_ffi / PartOpt / cpu / BothRev 0.000024976000000000003 s 0.000015853700006118743 s 1.58
hlo_ffi / IPartOpt / cpu / PreRev 0.000024332 s 0.000015407880000566364 s 1.58
hlo_ffi / IPartOpt / cpu / PostRev 0.000025257 s 0.000015501520019824967 s 1.63
hlo_ffi / IPartOpt / cpu / BothRev 0.000024576 s 0.00001576622002176009 s 1.56
hlo_ffi / DefOpt / cpu / PreRev 0.000024489 s 0.000015992219960025978 s 1.53
hlo_ffi / DefOpt / cpu / PostRev 0.000024592 s 0.000015379819997178858 s 1.60
hlo_ffi / DefOpt / cpu / BothRev 0.000025143 s 0.000015584500006298184 s 1.61
hlo_ffi / IDefOpt / cpu / PreRev 0.000024165 s 0.000015718799995738665 s 1.54
hlo_ffi / IDefOpt / cpu / PostRev 0.000025542 s 0.000015333300034399144 s 1.67
hlo_ffi / IDefOpt / cpu / BothRev 0.000025081 s 0.000015546340018772754 s 1.61
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / Primal 0.0011953959998209 s 0.0011389066000447 s 1.05
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / Primal 0.0009815054001592 s 0.0009614714001145 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / Primal 0.0010042806000456 s 0.0009669938000115 s 1.04
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / Primal 0.0009367898000164 s 0.0009036465999088 s 1.04
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / Primal 0.0009339360000012 s 0.0009796298000765 s 0.95
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / Primal 0.0009875813999315 s 0.0009802251999644 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / Primal 0.0009745836000547 s 0.0009544247999656 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / Forward 0.0028044923999914 s 0.0029485540000678 s 0.95
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / Forward 0.0024172855999495 s 0.0024281695999889 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / Forward 0.0023577881999699 s 0.0023517207999248 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / Forward 0.0023738597999908 s 0.0023108835999664 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / Forward 0.0023563141998238 s 0.0023455092001313 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / Forward 0.002318871800071 s 0.0025466422000135 s 0.91
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / Forward 0.0023425006000252 s 0.0022995338001237 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / PreRev 0.0064727188000688 s 0.0056316531999982 s 1.15
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / PostRev 0.0057140053999319 s 0.0058624548000807 s 0.97
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / BothRev 0.0054956565998509 s 0.0052789269999266 s 1.04
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / BothRev 0.0057205696000892 s 0.0053872159998718 s 1.06
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / PreRev 0.0052801766000811 s 0.0053421269999489 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / PostRev 0.0054021576000195 s 0.0050962933999471 s 1.06
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / BothRev 0.0053512106000198 s 0.0053552024000055 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / PreRev 0.005371247199946 s 0.0045724182000412 s 1.17
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / PostRev 0.0053587310001603 s 0.0057914356001674 s 0.93
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / BothRev 0.0052441319998251 s 0.0035167963999811 s 1.49
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / PreRev 0.0053551132001302 s 0.0053226364001602 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / PostRev 0.0054387301998758 s 0.003688061199864 s 1.47
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / BothRev 0.0052875077998578 s 0.0053972733999216 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / PreRev 0.005262680800115 s 0.004328001200065 s 1.22
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / PostRev 0.0057274835998214 s 0.0051587899999503 s 1.11
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / BothRev 0.0053420254001139 s 0.0035836637999636 s 1.49
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / PreRev 0.0054633201999422 s 0.0053208182000162 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / PostRev 0.005250640199938 s 0.0035166720000233 s 1.49
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / BothRev 0.0053575489998365 s 0.0053604224000082 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / Primal 0.0002823049999999 s 0.000273697 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cuda / Primal 0.0002814729999999 s 0.000272704 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cuda / Primal 0.000289985 s 0.000287073 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / Primal 0.000282817 s 0.00027264 s 1.04
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / Primal 0.000283458 s 0.000273408 s 1.04
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cuda / Primal 0.000290082 s 0.000287297 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cuda / Primal 0.000289409 s 0.000286977 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / Forward 0.000561795 s 0.000557409 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cuda / Forward 0.0005412829999999 s 0.000539138 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cuda / Forward 0.000563362 s 0.000557666 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / Forward 0.000562434 s 0.000557762 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / Forward 0.000560099 s 0.00055853 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cuda / Forward 0.0005610589999999 s 0.000558146 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cuda / Forward 0.000560994 s 0.000557857 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / PreRev 0.001052068 s 0.001022467 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / PostRev 0.001012515 s 0.000985283 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / BothRev 0.001048133 s 0.001019234 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cuda / BothRev 0.001008709 s 0.000979234 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cuda / PreRev 0.001035556 s 0.001006947 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cuda / PostRev 0.001059749 s 0.001029091 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cuda / BothRev 0.001033251 s 0.001006211 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / PreRev 0.001050627 s 0.001020899 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / PostRev 0.000997828 s 0.000970883 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / BothRev 0.001050372 s 0.001020387 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / PreRev 0.001046917 s 0.0010213779999999 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / PostRev 0.000999748 s 0.000972674 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / BothRev 0.001046404 s 0.001019586 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cuda / PreRev 0.001045124 s 0.001016707 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cuda / PostRev 0.000981988 s 0.0009542429999999 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cuda / BothRev 0.001046308 s 0.001016771 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cuda / PreRev 0.001045668 s 0.001017827 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cuda / PostRev 0.001044292 s 0.001017539 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cuda / BothRev 0.00104346 s 0.001015939 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / tpu / Primal 0.0001238037499999 s 0.0001283052499999 s 0.96
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / tpu / Primal 0.00012668075 s 0.0001240242499999 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / tpu / Primal 0.00015235625 s 0.0001577635 s 0.97
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / tpu / Primal 0.0001342985 s 0.000131264 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / tpu / Primal 0.00013077825 s 0.00013574875 s 0.96
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / tpu / Primal 0.000148273 s 0.00014559875 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / tpu / Primal 0.00015075525 s 0.000155781 s 0.97
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / tpu / Forward 0.00021207125 s 0.0002139735 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / tpu / Forward 0.00026096875 s 0.00026110975 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / tpu / Forward 0.0002122329999999 s 0.000220605 s 0.96
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / tpu / Forward 0.0002184522499999 s 0.0002136425 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / tpu / Forward 0.0002118834999999 s 0.0002162197499999 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / tpu / Forward 0.00021862325 s 0.0002181402499999 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / tpu / Forward 0.00021213825 s 0.00021671675 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / tpu / PreRev 0.00035492575 s 0.0003563254999999 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / tpu / PostRev 0.0002567149999999 s 0.000256164 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / tpu / BothRev 0.000354869 s 0.00035655275 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / tpu / BothRev 0.000257075 s 0.0002567605 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / tpu / PreRev 0.000354748 s 0.00035680425 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / tpu / PostRev 0.000290383 s 0.0002905205 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / tpu / BothRev 0.0003545955 s 0.0003565542499999 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / tpu / PreRev 0.00035615825 s 0.00035520425 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / tpu / PostRev 0.00027109475 s 0.0002721874999999 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / tpu / BothRev 0.00035567425 s 0.00035497025 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / tpu / PreRev 0.00035461925 s 0.00035670975 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / tpu / PostRev 0.000271627 s 0.00027145775 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / tpu / BothRev 0.00035482575 s 0.000356575 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / tpu / PreRev 0.0003580075 s 0.00035750575 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / tpu / PostRev 0.000283633 s 0.000283929 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / tpu / BothRev 0.00035795325 s 0.00035767725 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / tpu / PreRev 0.00035658175 s 0.00035885725 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / tpu / PostRev 0.0003012015 s 0.000301216 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / tpu / BothRev 0.0003569825 s 0.0003594085 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / Primal 0.001732606 s 0.0011389066000447 s 1.52
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / Primal 0.001738052 s 0.0009614714001145 s 1.81
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / Primal 0.001780738 s 0.0009669938000115 s 1.84
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / Primal 0.00157399 s 0.0009036465999088 s 1.74
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / Primal 0.001785713 s 0.0009796298000765 s 1.82
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / Primal 0.00145345 s 0.0009802251999644 s 1.48
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / Primal 0.001503969 s 0.0009544247999656 s 1.58
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / Forward 0.0049963329999999 s 0.0029485540000678 s 1.69
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / Forward 0.0048691279999999 s 0.0024281695999889 s 2.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / Forward 0.004780505 s 0.0023517207999248 s 2.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / Forward 0.004737439 s 0.0023108835999664 s 2.05
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / Forward 0.004715075 s 0.0023455092001313 s 2.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / Forward 0.005276434 s 0.0025466422000135 s 2.07
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / Forward 0.004892255 s 0.0022995338001237 s 2.13
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / PreRev 0.009624466 s 0.0056316531999982 s 1.71
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / PostRev 0.0078868669999999 s 0.0058624548000807 s 1.35
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / BothRev 0.007873229 s 0.0052789269999266 s 1.49
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / BothRev 0.008097408 s 0.0053872159998718 s 1.50
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / PreRev 0.008058141 s 0.0053421269999489 s 1.51
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / PostRev 0.007133829 s 0.0050962933999471 s 1.40
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / BothRev 0.008431578 s 0.0053552024000055 s 1.57
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / PreRev 0.009008489 s 0.0045724182000412 s 1.97
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / PostRev 0.008275801 s 0.0057914356001674 s 1.43
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / BothRev 0.008278406 s 0.0035167963999811 s 2.35
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / PreRev 0.0080889279999999 s 0.0053226364001602 s 1.52
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / PostRev 0.008323746 s 0.003688061199864 s 2.26
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / BothRev 0.007612168 s 0.0053972733999216 s 1.41
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / PreRev 0.008378372 s 0.004328001200065 s 1.94
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / PostRev 0.007081088 s 0.0051587899999503 s 1.37
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / BothRev 0.0084766069999999 s 0.0035836637999636 s 2.37
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / PreRev 0.007700518 s 0.0053208182000162 s 1.45
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / PostRev 0.008528401 s 0.0035166720000233 s 2.43
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / BothRev 0.007745165 s 0.0053604224000082 s 1.44
scatter_sum / JaXPipe / cpu / Primal 0.00000949315994148492 s 0.00001004172002467385 s 0.95
scatter_sum / Jax / cpu / Primal 0.000008826519997455761 s 0.000009407980005562422 s 0.94
scatter_sum / HLOOpt / cpu / Primal 0.000012542819995360332 s 0.000012474039995140628 s 1.01
scatter_sum / PartOpt / cpu / Primal 0.000009376199996040668 s 0.000008274519987025996 s 1.13
scatter_sum / IPartOpt / cpu / Primal 0.000008971720026238472 s 0.000008602779998909683 s 1.04
scatter_sum / DefOpt / cpu / Primal 0.000009283240024160475 s 0.000008273100002043066 s 1.12
scatter_sum / IDefOpt / cpu / Primal 0.000009005260026242467 s 0.000008372360016437596 s 1.08
scatter_sum / JaXPipe / cpu / Forward 0.000013522199988074134 s 0.000013005980044908935 s 1.04
scatter_sum / Jax / cpu / Forward 0.000012886879994766789 s 0.00001275019997592608 s 1.01
scatter_sum / HLOOpt / cpu / Forward 0.000018923480001831195 s 0.00001813943997149181 s 1.04
scatter_sum / PartOpt / cpu / Forward 0.000013716299990846892 s 0.00001831355997637729 s 0.75
scatter_sum / IPartOpt / cpu / Forward 0.000013248620034573832 s 0.000012779299959220224 s 1.04
scatter_sum / DefOpt / cpu / Forward 0.00001888731996587012 s 0.000018687720003072174 s 1.01
scatter_sum / IDefOpt / cpu / Forward 0.000013605239992102725 s 0.00001276049996704387 s 1.07
scatter_sum / JaXPipe / cpu / PreRev 0.00001469114001338312 s 0.000013078879983368096 s 1.12
scatter_sum / JaXPipe / cpu / PostRev 0.000013005079972572276 s 0.000012596560036399751 s 1.03
scatter_sum / JaXPipe / cpu / BothRev 0.000018056739963867582 s 0.000012585200020112095 s 1.43
scatter_sum / Jax / cpu / BothRev 0.000014025659993421868 s 0.00001235963996805367 s 1.13
scatter_sum / HLOOpt / cpu / PreRev 0.000013765300009254132 s 0.00001306199997998192 s 1.05
scatter_sum / HLOOpt / cpu / PostRev 0.00001335312001174316 s 0.00001713214001938468 s 0.78
scatter_sum / HLOOpt / cpu / BothRev 0.000015180700002019876 s 0.00002021306001552148 s 0.75
scatter_sum / PartOpt / cpu / PreRev 0.000013215700018918142 s 0.000012858079999205077 s 1.03
scatter_sum / PartOpt / cpu / PostRev 0.000013610639971375347 s 0.000013021439999647554 s 1.05
scatter_sum / PartOpt / cpu / BothRev 0.000013090840029690298 s 0.00001270738001949212 s 1.03
scatter_sum / IPartOpt / cpu / PreRev 0.000019554979999156782 s 0.000013217800005804748 s 1.48
scatter_sum / IPartOpt / cpu / PostRev 0.000012978799968550448 s 0.000013108480006849276 s 0.99
scatter_sum / IPartOpt / cpu / BothRev 0.000013130760034982812 s 0.000012461719979910412 s 1.05
scatter_sum / DefOpt / cpu / PreRev 0.000013230800013843692 s 0.00001264913998056727 s 1.05
scatter_sum / DefOpt / cpu / PostRev 0.00001309726000727096 s 0.0000127456999689457 s 1.03
scatter_sum / DefOpt / cpu / BothRev 0.000013270300014482927 s 0.0000126665800507908 s 1.05
scatter_sum / IDefOpt / cpu / PreRev 0.000013460339978337288 s 0.0000127709199932724 s 1.05
scatter_sum / IDefOpt / cpu / PostRev 0.000012928439973620698 s 0.00001324225998359907 s 0.98
scatter_sum / IDefOpt / cpu / BothRev 0.0000130624000030366 s 0.000013151660004950828 s 0.99
scatter_sum / JaXPipe / cuda / Primal 0.00001024 s 0.000010464 s 0.98
scatter_sum / Jax / cuda / Primal 0.000009984 s 0.000010464 s 0.95
scatter_sum / HLOOpt / cuda / Primal 0.000010112 s 0.00001024 s 0.99
scatter_sum / PartOpt / cuda / Primal 0.000010144 s 0.000010368 s 0.98
scatter_sum / IPartOpt / cuda / Primal 0.000010176 s 0.000010048 s 1.01
scatter_sum / DefOpt / cuda / Primal 0.000009631 s 0.000010176 s 0.95
scatter_sum / IDefOpt / cuda / Primal 0.00000992 s 0.000010497 s 0.95
scatter_sum / JaXPipe / cuda / Forward 0.000017633 s 0.000017472 s 1.01
scatter_sum / Jax / cuda / Forward 0.000016927999999999998 s 0.00001744 s 0.97
scatter_sum / HLOOpt / cuda / Forward 0.00001696 s 0.000017408 s 0.97
scatter_sum / PartOpt / cuda / Forward 0.000017503999999999997 s 0.000017824 s 0.98
scatter_sum / IPartOpt / cuda / Forward 0.000017472 s 0.000017569 s 0.99
scatter_sum / DefOpt / cuda / Forward 0.000017281000000000003 s 0.000017984 s 0.96
scatter_sum / IDefOpt / cuda / Forward 0.000017375999999999998 s 0.000018272 s 0.95
scatter_sum / JaXPipe / cuda / PreRev 0.000017696 s 0.000018112 s 0.98
scatter_sum / JaXPipe / cuda / PostRev 0.00001728 s 0.000018144 s 0.95
scatter_sum / JaXPipe / cuda / BothRev 0.000017184 s 0.00002016 s 0.85
scatter_sum / Jax / cuda / BothRev 0.000016992 s 0.000019776 s 0.86
scatter_sum / HLOOpt / cuda / PreRev 0.000018432 s 0.000020128 s 0.92
scatter_sum / HLOOpt / cuda / PostRev 0.000017919999999999998 s 0.000018207 s 0.98
scatter_sum / HLOOpt / cuda / BothRev 0.000017792 s 0.000017919999999999998 s 0.99
scatter_sum / PartOpt / cuda / PreRev 0.00001824 s 0.000018047 s 1.01
scatter_sum / PartOpt / cuda / PostRev 0.00001744 s 0.000018144 s 0.96
scatter_sum / PartOpt / cuda / BothRev 0.000017921 s 0.000018144 s 0.99
scatter_sum / IPartOpt / cuda / PreRev 0.000016704 s 0.000018208 s 0.92
scatter_sum / IPartOpt / cuda / PostRev 0.000017408 s 0.00001728 s 1.01
scatter_sum / IPartOpt / cuda / BothRev 0.00001712 s 0.000017696 s 0.97
scatter_sum / DefOpt / cuda / PreRev 0.0000176 s 0.00002016 s 0.87
scatter_sum / DefOpt / cuda / PostRev 0.000017247999999999998 s 0.000017793 s 0.97
scatter_sum / DefOpt / cuda / BothRev 0.0000176 s 0.000018016 s 0.98
scatter_sum / IDefOpt / cuda / PreRev 0.000017760000000000003 s 0.000018112 s 0.98
scatter_sum / IDefOpt / cuda / PostRev 0.000017375000000000002 s 0.000018016 s 0.96
scatter_sum / IDefOpt / cuda / BothRev 0.000017056 s 0.000018176 s 0.94
scatter_sum / JaXPipe / tpu / Primal 0.0000013502249999999998 s 0.000001344175 s 1.00
scatter_sum / Jax / tpu / Primal 0.0000014139 s 0.0000013536000000000005 s 1.04
scatter_sum / HLOOpt / tpu / Primal 0.00000135965 s 0.0000013533 s 1.00
scatter_sum / PartOpt / tpu / Primal 0.0000014135 s 0.000001353125 s 1.04
scatter_sum / IPartOpt / tpu / Primal 0.000001359825 s 0.000001353625 s 1.00
scatter_sum / DefOpt / tpu / Primal 0.0000014146 s 0.000001353425 s 1.05
scatter_sum / IDefOpt / tpu / Primal 0.000001361075 s 0.000001353475 s 1.01
scatter_sum / JaXPipe / tpu / Forward 0.000002716 s 0.0000026949 s 1.01
scatter_sum / Jax / tpu / Forward 0.000002732875 s 0.000002733775 s 1.00
scatter_sum / HLOOpt / tpu / Forward 0.000002712 s 0.000002692325 s 1.01
scatter_sum / PartOpt / tpu / Forward 0.0000027028500000000003 s 0.00000270215 s 1.00
scatter_sum / IPartOpt / tpu / Forward 0.00000271185 s 0.0000026971250000000005 s 1.01
scatter_sum / DefOpt / tpu / Forward 0.0000027019500000000004 s 0.000002701375 s 1.00
scatter_sum / IDefOpt / tpu / Forward 0.000002713925 s 0.00000268995 s 1.01
scatter_sum / JaXPipe / tpu / PreRev 0.0000027014 s 0.000002696 s 1.00
scatter_sum / JaXPipe / tpu / PostRev 0.000002699775 s 0.0000026889 s 1.00
scatter_sum / JaXPipe / tpu / BothRev 0.00000271115 s 0.0000027099 s 1.00
scatter_sum / Jax / tpu / BothRev 0.000002747325 s 0.000002743375 s 1.00
scatter_sum / HLOOpt / tpu / PreRev 0.0000027237250000000003 s 0.000002711775 s 1.00
scatter_sum / HLOOpt / tpu / PostRev 0.000002754175 s 0.0000027473 s 1.00
scatter_sum / HLOOpt / tpu / BothRev 0.000002713425 s 0.000002710375 s 1.00
scatter_sum / PartOpt / tpu / PreRev 0.000002750875 s 0.000002751025 s 1.00
scatter_sum / PartOpt / tpu / PostRev 0.000002712175 s 0.000002709625 s 1.00
scatter_sum / PartOpt / tpu / BothRev 0.0000027507 s 0.00000274645 s 1.00
scatter_sum / IPartOpt / tpu / PreRev 0.0000027151 s 0.000002707175 s 1.00
scatter_sum / IPartOpt / tpu / PostRev 0.0000027484250000000003 s 0.00000274945 s 1.00
scatter_sum / IPartOpt / tpu / BothRev 0.0000027158 s 0.0000027067000000000003 s 1.00
scatter_sum / DefOpt / tpu / PreRev 0.0000027489000000000005 s 0.000002742125 s 1.00
scatter_sum / DefOpt / tpu / PostRev 0.0000027131 s 0.000002714125 s 1.00
scatter_sum / DefOpt / tpu / BothRev 0.000002751925 s 0.00000274985 s 1.00
scatter_sum / IDefOpt / tpu / PreRev 0.00000272955 s 0.00000270505 s 1.01
scatter_sum / IDefOpt / tpu / PostRev 0.0000027575250000000003 s 0.0000027458 s 1.00
scatter_sum / IDefOpt / tpu / BothRev 0.000002715875 s 0.000002704625 s 1.00
scatter_sum / JaXPipe / cpu / Primal 0.000015589 s 0.00001004172002467385 s 1.55
scatter_sum / Jax / cpu / Primal 0.000015233 s 0.000009407980005562422 s 1.62
scatter_sum / HLOOpt / cpu / Primal 0.000015565 s 0.000012474039995140628 s 1.25
scatter_sum / PartOpt / cpu / Primal 0.000015252 s 0.000008274519987025996 s 1.84
scatter_sum / IPartOpt / cpu / Primal 0.000015256 s 0.000008602779998909683 s 1.77
scatter_sum / DefOpt / cpu / Primal 0.000015624 s 0.000008273100002043066 s 1.89
scatter_sum / IDefOpt / cpu / Primal 0.000015498 s 0.000008372360016437596 s 1.85
scatter_sum / JaXPipe / cpu / Forward 0.000022132 s 0.000013005980044908935 s 1.70
scatter_sum / Jax / cpu / Forward 0.000022661 s 0.00001275019997592608 s 1.78
scatter_sum / HLOOpt / cpu / Forward 0.00002167 s 0.00001813943997149181 s 1.19
scatter_sum / PartOpt / cpu / Forward 0.0000223 s 0.00001831355997637729 s 1.22
scatter_sum / IPartOpt / cpu / Forward 0.000022004 s 0.000012779299959220224 s 1.72
scatter_sum / DefOpt / cpu / Forward 0.000022665 s 0.000018687720003072174 s 1.21
scatter_sum / IDefOpt / cpu / Forward 0.000021991 s 0.00001276049996704387 s 1.72
scatter_sum / JaXPipe / cpu / PreRev 0.000022558000000000003 s 0.000013078879983368096 s 1.72
scatter_sum / JaXPipe / cpu / PostRev 0.000022627 s 0.000012596560036399751 s 1.80
scatter_sum / JaXPipe / cpu / BothRev 0.000023478 s 0.000012585200020112095 s 1.87
scatter_sum / Jax / cpu / BothRev 0.000022301 s 0.00001235963996805367 s 1.80
scatter_sum / HLOOpt / cpu / PreRev 0.000022329000000000003 s 0.00001306199997998192 s 1.71
scatter_sum / HLOOpt / cpu / PostRev 0.000022872 s 0.00001713214001938468 s 1.34
scatter_sum / HLOOpt / cpu / BothRev 0.000022056 s 0.00002021306001552148 s 1.09
scatter_sum / PartOpt / cpu / PreRev 0.000022334 s 0.000012858079999205077 s 1.74
scatter_sum / PartOpt / cpu / PostRev 0.000022888 s 0.000013021439999647554 s 1.76
scatter_sum / PartOpt / cpu / BothRev 0.000023351 s 0.00001270738001949212 s 1.84
scatter_sum / IPartOpt / cpu / PreRev 0.000022364 s 0.000013217800005804748 s 1.69
scatter_sum / IPartOpt / cpu / PostRev 0.000022924 s 0.000013108480006849276 s 1.75
scatter_sum / IPartOpt / cpu / BothRev 0.000022484 s 0.000012461719979910412 s 1.80
scatter_sum / DefOpt / cpu / PreRev 0.000021639 s 0.00001264913998056727 s 1.71
scatter_sum / DefOpt / cpu / PostRev 0.000022054 s 0.0000127456999689457 s 1.73
scatter_sum / DefOpt / cpu / BothRev 0.000022435 s 0.0000126665800507908 s 1.77
scatter_sum / IDefOpt / cpu / PreRev 0.000021843 s 0.0000127709199932724 s 1.71
scatter_sum / IDefOpt / cpu / PostRev 0.000022075 s 0.00001324225998359907 s 1.67
scatter_sum / IDefOpt / cpu / BothRev 0.000022629 s 0.000013151660004950828 s 1.72
slicing / JaXPipe / cpu / Primal 0.000007514640001318184 s 0.000007570659972770954 s 0.99
slicing / Jax / cpu / Primal 0.000006591599985767971 s 0.000006614699987039785 s 1.00
slicing / HLOOpt / cpu / Primal 0.00001061686001776252 s 0.000010874300032810425 s 0.98
slicing / PartOpt / cpu / Primal 0.000006807019999541808 s 0.000006771719990865677 s 1.01
slicing / IPartOpt / cpu / Primal 0.00000701588000083575 s 0.000007026199973552139 s 1.00
slicing / DefOpt / cpu / Primal 0.00001120758001889044 s 0.000011422439993111766 s 0.98
slicing / IDefOpt / cpu / Primal 0.00000707851995684905 s 0.000006992520020503435 s 1.01
slicing / JaXPipe / cpu / Forward 0.00001066877997800475 s 0.000010335900005884469 s 1.03
slicing / Jax / cpu / Forward 0.000011374000032446928 s 0.0000114975600172329 s 0.99
slicing / HLOOpt / cpu / Forward 0.000014966759981689392 s 0.000014506579973385669 s 1.03
slicing / PartOpt / cpu / Forward 0.000014820739988863353 s 0.000014642859960076748 s 1.01
slicing / IPartOpt / cpu / Forward 0.000009968480026145698 s 0.000010182800015172688 s 0.98
slicing / DefOpt / cpu / Forward 0.000014738640002178727 s 0.000014841600022919013 s 0.99
slicing / IDefOpt / cpu / Forward 0.00001019878004626662 s 0.0000099242199848959 s 1.03
slicing / JaXPipe / cpu / PreRev 0.00001091387994165416 s 0.000011109560018667253 s 0.98
slicing / JaXPipe / cpu / PostRev 0.00001123116001508606 s 0.00001135981993684254 s 0.99
slicing / JaXPipe / cpu / BothRev 0.000010773040003186908 s 0.000014937100022507366 s 0.72
slicing / Jax / cpu / BothRev 0.000010771900015242864 s 0.000011023059978469973 s 0.98
slicing / HLOOpt / cpu / PreRev 0.000010741680043793166 s 0.000011034760009351884 s 0.97
slicing / HLOOpt / cpu / PostRev 0.000011237980015721404 s 0.00001136896004936716 s 0.99
slicing / HLOOpt / cpu / BothRev 0.000012705539993476123 s 0.000012524459989435857 s 1.01
slicing / PartOpt / cpu / PreRev 0.000010796200022014092 s 0.000010715039989008802 s 1.01
slicing / PartOpt / cpu / PostRev 0.000011201520028407683 s 0.000011674460020003608 s 0.96
slicing / PartOpt / cpu / BothRev 0.000011081520024163182 s 0.000010953879991575375 s 1.01
slicing / IPartOpt / cpu / PreRev 0.00001566471999467467 s 0.00001082678002603643 s 1.45
slicing / IPartOpt / cpu / PostRev 0.000011564439982976184 s 0.000011218379995625584 s 1.03
slicing / IPartOpt / cpu / BothRev 0.000010703299994929694 s 0.00001056629998856806 s 1.01
slicing / DefOpt / cpu / PreRev 0.00001061505997313361 s 0.00001046529998347978 s 1.01
slicing / DefOpt / cpu / PostRev 0.00001073012003871554 s 0.000011711559973264229 s 0.92
slicing / DefOpt / cpu / BothRev 0.000010802739971040864 s 0.000010657140037437783 s 1.01
slicing / IDefOpt / cpu / PreRev 0.000011105700014013563 s 0.00001072921995728393 s 1.04
slicing / IDefOpt / cpu / PostRev 0.00001124074001381814 s 0.000011978920001638471 s 0.94
slicing / IDefOpt / cpu / BothRev 0.000010790760034069536 s 0.000010620360008033458 s 1.02
slicing / JaXPipe / cuda / Primal 0.0000019200000000000003 s 0.000001887 s 1.02
slicing / Jax / cuda / Primal 0.000001888 s 0.000001887 s 1.00
slicing / HLOOpt / cuda / Primal 0.000001888 s 0.000001887 s 1.00
slicing / PartOpt / cuda / Primal 0.000001888 s 0.000001887 s 1.00
slicing / IPartOpt / cuda / Primal 0.0000019200000000000003 s 0.000001887 s 1.02
slicing / DefOpt / cuda / Primal 0.0000019200000000000003 s 0.000001887 s 1.02
slicing / IDefOpt / cuda / Primal 0.000001888 s 0.000001887 s 1.00
slicing / JaXPipe / cuda / Forward 0.000009537 s 0.000010368 s 0.92
slicing / Jax / cuda / Forward 0.000010016 s 0.000010144 s 0.99
slicing / HLOOpt / cuda / Forward 0.000010304 s 0.00000976 s 1.06
slicing / PartOpt / cuda / Forward 0.000010144 s 0.000010368 s 0.98
slicing / IPartOpt / cuda / Forward 0.000009984 s 0.000010177 s 0.98
slicing / DefOpt / cuda / Forward 0.00000976 s 0.000011681 s 0.84
slicing / IDefOpt / cuda / Forward 0.00000992 s 0.000010016 s 0.99
slicing / JaXPipe / cuda / PreRev 0.00001024 s 0.000010656 s 0.96
slicing / JaXPipe / cuda / PostRev 0.000010144 s 0.000010176 s 1.00
slicing / JaXPipe / cuda / BothRev 0.000010175 s 0.000010848 s 0.94
slicing / Jax / cuda / BothRev 0.0000104 s 0.000011616 s 0.90
slicing / HLOOpt / cuda / PreRev 0.00001024 s 0.000010625 s 0.96
slicing / HLOOpt / cuda / PostRev 0.000010208 s 0.000010656 s 0.96
slicing / HLOOpt / cuda / BothRev 0.000010432 s 0.000011616 s 0.90
slicing / PartOpt / cuda / PreRev 0.000010048 s 0.000010593 s 0.95
slicing / PartOpt / cuda / PostRev 0.000010176 s 0.000009792 s 1.04
slicing / PartOpt / cuda / BothRev 0.000009889 s 0.000010816 s 0.91
slicing / IPartOpt / cuda / PreRev 0.000010464 s 0.00001056 s 0.99
slicing / IPartOpt / cuda / PostRev 0.000010368 s 0.000011744 s 0.88
slicing / IPartOpt / cuda / BothRev 0.000010208 s 0.00001184 s 0.86
slicing / DefOpt / cuda / PreRev 0.000010432 s 0.000010687 s 0.98
slicing / DefOpt / cuda / PostRev 0.000010432 s 0.000009536 s 1.09
slicing / DefOpt / cuda / BothRev 0.000010144 s 0.000010529 s 0.96
slicing / IDefOpt / cuda / PreRev 0.000010144 s 0.00001056 s 0.96
slicing / IDefOpt / cuda / PostRev 0.00001056 s 0.000010496 s 1.01
slicing / IDefOpt / cuda / BothRev 0.000010144 s 0.00001056 s 0.96
slicing / JaXPipe / tpu / Primal 9.65225e-7 s 9.605e-7 s 1.00
slicing / Jax / tpu / Primal 9.67075e-7 s 9.55225e-7 s 1.01
slicing / HLOOpt / tpu / Primal 9.73275e-7 s 9.67975e-7 s 1.01
slicing / PartOpt / tpu / Primal 9.70225e-7 s 9.6795e-7 s 1.00
slicing / IPartOpt / tpu / Primal 9.6435e-7 s 9.58875e-7 s 1.01
slicing / DefOpt / tpu / Primal 9.728e-7 s 9.566e-7 s 1.02
slicing / IDefOpt / tpu / Primal 9.64975e-7 s 9.67125e-7 s 1.00
slicing / JaXPipe / tpu / Forward 0.0000014137 s 0.00000140265 s 1.01
slicing / Jax / tpu / Forward 0.0000014303250000000002 s 0.000001409225 s 1.01
slicing / HLOOpt / tpu / Forward 0.00000151635 s 0.00000151315 s 1.00
slicing / PartOpt / tpu / Forward 0.0000014487 s 0.0000014359 s 1.01
slicing / IPartOpt / tpu / Forward 0.0000015222 s 0.0000015109 s 1.01
slicing / DefOpt / tpu / Forward 0.000001436225 s 0.00000142585 s 1.01
slicing / IDefOpt / tpu / Forward 0.00000152095 s 0.0000015207 s 1.00
slicing / JaXPipe / tpu / PreRev 0.0000023749 s 0.000002342725 s 1.01
slicing / JaXPipe / tpu / PostRev 0.0000025283 s 0.000002510175 s 1.01
slicing / JaXPipe / tpu / BothRev 0.0000024048 s 0.00000235005 s 1.02
slicing / Jax / tpu / BothRev 0.0000025350250000000003 s 0.0000025342750000000004 s 1.00
slicing / HLOOpt / tpu / PreRev 0.000002399825 s 0.000002348625 s 1.02
slicing / HLOOpt / tpu / PostRev 0.000002539075 s 0.0000025253250000000003 s 1.01
slicing / HLOOpt / tpu / BothRev 0.000002395675 s 0.0000023492 s 1.02
slicing / PartOpt / tpu / PreRev 0.00000253705 s 0.0000025327000000000003 s 1.00
slicing / PartOpt / tpu / PostRev 0.000002393925 s 0.0000023542750000000004 s 1.02
slicing / PartOpt / tpu / BothRev 0.000002544525 s 0.0000025369750000000004 s 1.00
slicing / IPartOpt / tpu / PreRev 0.000002401225 s 0.00000236045 s 1.02
slicing / IPartOpt / tpu / PostRev 0.000002560225 s 0.000002527725 s 1.01
slicing / IPartOpt / tpu / BothRev 0.000002394875 s 0.000002354475 s 1.02
slicing / DefOpt / tpu / PreRev 0.00000254575 s 0.000002526675 s 1.01
slicing / DefOpt / tpu / PostRev 0.0000023950250000000004 s 0.0000023534 s 1.02
slicing / DefOpt / tpu / BothRev 0.00000254765 s 0.0000025356000000000003 s 1.00
slicing / IDefOpt / tpu / PreRev 0.000002400525 s 0.000002347575 s 1.02
slicing / IDefOpt / tpu / PostRev 0.000002542925 s 0.00000253375 s 1.00
slicing / IDefOpt / tpu / BothRev 0.00000240295 s 0.0000023546000000000003 s 1.02
slicing / JaXPipe / cpu / Primal 0.000012673 s 0.000007570659972770954 s 1.67
slicing / Jax / cpu / Primal 0.000012423 s 0.000006614699987039785 s 1.88
slicing / HLOOpt / cpu / Primal 0.000012642 s 0.000010874300032810425 s 1.16
slicing / PartOpt / cpu / Primal 0.000012236 s 0.000006771719990865677 s 1.81
slicing / IPartOpt / cpu / Primal 0.000012085 s 0.000007026199973552139 s 1.72
slicing / DefOpt / cpu / Primal 0.000012124 s 0.000011422439993111766 s 1.06
slicing / IDefOpt / cpu / Primal 0.000012244 s 0.000006992520020503435 s 1.75
slicing / JaXPipe / cpu / Forward 0.000017118 s 0.000010335900005884469 s 1.66
slicing / Jax / cpu / Forward 0.000016284 s 0.0000114975600172329 s 1.42
slicing / HLOOpt / cpu / Forward 0.000016219000000000002 s 0.000014506579973385669 s 1.12
slicing / PartOpt / cpu / Forward 0.000016471 s 0.000014642859960076748 s 1.12
slicing / IPartOpt / cpu / Forward 0.000016195 s 0.000010182800015172688 s 1.59
slicing / DefOpt / cpu / Forward 0.000016414000000000002 s 0.000014841600022919013 s 1.11
slicing / IDefOpt / cpu / Forward 0.000016733 s 0.0000099242199848959 s 1.69
slicing / JaXPipe / cpu / PreRev 0.000017255 s 0.000011109560018667253 s 1.55
slicing / JaXPipe / cpu / PostRev 0.000017754 s 0.00001135981993684254 s 1.56
slicing / JaXPipe / cpu / BothRev 0.000017760000000000003 s 0.000014937100022507366 s 1.19
slicing / Jax / cpu / BothRev 0.00001732 s 0.000011023059978469973 s 1.57
slicing / HLOOpt / cpu / PreRev 0.000017246 s 0.000011034760009351884 s 1.56
slicing / HLOOpt / cpu / PostRev 0.000017339 s 0.00001136896004936716 s 1.53
slicing / HLOOpt / cpu / BothRev 0.000017422 s 0.000012524459989435857 s 1.39
slicing / PartOpt / cpu / PreRev 0.000017017 s 0.000010715039989008802 s 1.59
slicing / PartOpt / cpu / PostRev 0.00001726 s 0.000011674460020003608 s 1.48
slicing / PartOpt / cpu / BothRev 0.000018018 s 0.000010953879991575375 s 1.64
slicing / IPartOpt / cpu / PreRev 0.000017229 s 0.00001082678002603643 s 1.59
slicing / IPartOpt / cpu / PostRev 0.000017386 s 0.000011218379995625584 s 1.55
slicing / IPartOpt / cpu / BothRev 0.000017715000000000002 s 0.00001056629998856806 s 1.68
slicing / DefOpt / cpu / PreRev 0.000017316 s 0.00001046529998347978 s 1.65
slicing / DefOpt / cpu / PostRev 0.00001782 s 0.000011711559973264229 s 1.52
slicing / DefOpt / cpu / BothRev 0.000017619 s 0.000010657140037437783 s 1.65
slicing / IDefOpt / cpu / PreRev 0.000017214 s 0.00001072921995728393 s 1.60
slicing / IDefOpt / cpu / PostRev 0.000017641 s 0.000011978920001638471 s 1.47
slicing / IDefOpt / cpu / BothRev 0.000017281999999999998 s 0.000010620360008033458 s 1.63
sum / JaXPipe / cpu / Primal 0.000009729680004966211 s 0.00000879642001564207 s 1.11
sum / Jax / cpu / Primal 0.000008456140021735336 s 0.000007896020024418249 s 1.07
sum / HLOOpt / cpu / Primal 0.00001385242002470477 s 0.00001246118001290597 s 1.11
sum / PartOpt / cpu / Primal 0.000008547140014343313 s 0.000008450039977105917 s 1.01
sum / IPartOpt / cpu / Primal 0.000009120199993049029 s 0.000008476499997414067 s 1.08
sum / DefOpt / cpu / Primal 0.000009576680022291838 s 0.000012441299986676311 s 0.77
sum / IDefOpt / cpu / Primal 0.000009178140016956604 s 0.000008302099977299805 s 1.11
sum / JaXPipe / cpu / Forward 0.000012653739959205269 s 0.000012332720007179888 s 1.03
sum / Jax / cpu / Forward 0.000013194980010666767 s 0.000012614460001714178 s 1.05
sum / HLOOpt / cpu / Forward 0.00001785802004633297 s 0.00001719820002108463 s 1.04
sum / PartOpt / cpu / Forward 0.00001885453998511366 s 0.000016864640010680886 s 1.12
sum / IPartOpt / cpu / Forward 0.000012841939997088047 s 0.000012690599969573669 s 1.01
sum / DefOpt / cpu / Forward 0.00001870641995083133 s 0.000017216219976035064 s 1.09
sum / IDefOpt / cpu / Forward 0.000013104660019962468 s 0.000012815299996873363 s 1.02
sum / JaXPipe / cpu / PreRev 0.00001216808004755876 s 0.00001255094000043755 s 0.97
sum / JaXPipe / cpu / PostRev 0.000012265659979675549 s 0.00001229197998327436 s 1.00
sum / JaXPipe / cpu / BothRev 0.000012129539964007563 s 0.00001570249994074402 s 0.77
sum / Jax / cpu / BothRev 0.00001246285998604435 s 0.00001219255997966684 s 1.02
sum / HLOOpt / cpu / PreRev 0.000011506940036269952 s 0.00001153561995124619 s 1.00
sum / HLOOpt / cpu / PostRev 0.000016316680012096184 s 0.00001561490000312915 s 1.04
sum / HLOOpt / cpu / BothRev 0.00001367054000183998 s 0.000013388559946179155 s 1.02
sum / PartOpt / cpu / PreRev 0.000011884859995916486 s 0.00001177519997327181 s 1.01
sum / PartOpt / cpu / PostRev 0.000012069500025972956 s 0.000011917339970750615 s 1.01
sum / PartOpt / cpu / BothRev 0.000011923419951926918 s 0.000011633939984676544 s 1.02
sum / IPartOpt / cpu / PreRev 0.000012002940038655652 s 0.000015914320001684246 s 0.75
sum / IPartOpt / cpu / PostRev 0.000011545660026968109 s 0.00001144826000199828 s 1.01
sum / IPartOpt / cpu / BothRev 0.000011768939993999084 s 0.000011945719979848943 s 0.99
sum / DefOpt / cpu / PreRev 0.000011600059951888398 s 0.0000116967799840495 s 0.99
sum / DefOpt / cpu / PostRev 0.00001209818001370877 s 0.000012046660003761644 s 1.00
sum / DefOpt / cpu / BothRev 0.000011482660047477111 s 0.00001183411999591044 s 0.97
sum / IDefOpt / cpu / PreRev 0.000011267080008110498 s 0.000011886359998243278 s 0.95
sum / IDefOpt / cpu / PostRev 0.000011605340005189644 s 0.000012097099952370627 s 0.96
sum / IDefOpt / cpu / BothRev 0.000011723140023605084 s 0.000011305200005153892 s 1.04
sum / JaXPipe / cuda / Primal 0.00000208 s 0.000002047 s 1.02
sum / Jax / cuda / Primal 0.00000208 s 0.000002048 s 1.02
sum / HLOOpt / cuda / Primal 0.00000208 s 0.000002048 s 1.02
sum / PartOpt / cuda / Primal 0.00000208 s 0.000002047 s 1.02
sum / IPartOpt / cuda / Primal 0.00000208 s 0.000002047 s 1.02
sum / DefOpt / cuda / Primal 0.00000208 s 0.000002047 s 1.02
sum / IDefOpt / cuda / Primal 0.00000208 s 0.000002048 s 1.02
sum / JaXPipe / cuda / Forward 0.000010144 s 0.00001072 s 0.95
sum / Jax / cuda / Forward 0.00001072 s 0.00001056 s 1.02
sum / HLOOpt / cuda / Forward 0.000015392 s 0.000010431 s 1.48
sum / PartOpt / cuda / Forward 0.000010176 s 0.000010912 s 0.93
sum / IPartOpt / cuda / Forward 0.000010432 s 0.000010368 s 1.01
sum / DefOpt / cuda / Forward 0.000010689 s 0.00001072 s 1.00
sum / IDefOpt / cuda / Forward 0.000010528 s 0.000010367 s 1.02
sum / JaXPipe / cuda / PreRev 0.00001008 s 0.000009951 s 1.01
sum / JaXPipe / cuda / PostRev 0.000010209 s 0.000009887 s 1.03
sum / JaXPipe / cuda / BothRev 0.000010048 s 0.000010176 s 0.99
sum / Jax / cuda / BothRev 0.000010111 s 0.0000104 s 0.97
sum / HLOOpt / cuda / PreRev 0.00000976 s 0.000010335 s 0.94
sum / HLOOpt / cuda / PostRev 0.000009888 s 0.000011712 s 0.84
sum / HLOOpt / cuda / BothRev 0.00001008 s 0.000011169 s 0.90
sum / PartOpt / cuda / PreRev 0.000010112 s 0.000015966999999999998 s 0.63
sum / PartOpt / cuda / PostRev 0.000009887 s 0.000010592 s 0.93
sum / PartOpt / cuda / BothRev 0.00001008 s 0.000010496 s 0.96
sum / IPartOpt / cuda / PreRev 0.000010112 s 0.0000104 s 0.97
sum / IPartOpt / cuda / PostRev 0.000010144 s 0.000010784 s 0.94
sum / IPartOpt / cuda / BothRev 0.00000992 s 0.00001072 s 0.93
sum / DefOpt / cuda / PreRev 0.00001008 s 0.000010528 s 0.96
sum / DefOpt / cuda / PostRev 0.000009823 s 0.000010944 s 0.90
sum / DefOpt / cuda / BothRev 0.000010176 s 0.000010464 s 0.97
sum / IDefOpt / cuda / PreRev 0.000009952 s 0.000010528 s 0.95
sum / IDefOpt / cuda / PostRev 0.000010176 s 0.000010464 s 0.97
sum / IDefOpt / cuda / BothRev 0.00001008 s 0.000010368 s 0.97
sum / JaXPipe / tpu / Primal 5.105e-7 s 5.0325e-7 s 1.01
sum / Jax / tpu / Primal 5.57125e-7 s 5.57325e-7 s 1.00
sum / HLOOpt / tpu / Primal 5.2015e-7 s 5.1325e-7 s 1.01
sum / PartOpt / tpu / Primal 5.574500000000001e-7 s 5.574e-7 s 1.00
sum / IPartOpt / tpu / Primal 5.20325e-7 s 5.131e-7 s 1.01
sum / DefOpt / tpu / Primal 5.571e-7 s 5.56925e-7 s 1.00
sum / IDefOpt / tpu / Primal 5.207e-7 s 5.131750000000001e-7 s 1.01
sum / JaXPipe / tpu / Forward 0.00000155195 s 0.0000015524499999999998 s 1.00
sum / Jax / tpu / Forward 0.00000150055 s 0.000001492675 s 1.01
sum / HLOOpt / tpu / Forward 0.000001537525 s 0.0000015325 s 1.00
sum / PartOpt / tpu / Forward 0.0000014982500000000002 s 0.000001493725 s 1.00
sum / IPartOpt / tpu / Forward 0.00000153285 s 0.0000015295 s 1.00
sum / DefOpt / tpu / Forward 0.000001500375 s 0.0000014891 s 1.01
sum / IDefOpt / tpu / Forward 0.000001536325 s 0.0000015311749999999998 s 1.00
sum / JaXPipe / tpu / PreRev 9.99425e-7 s 0.000001003075 s 1.00
sum / JaXPipe / tpu / PostRev 0.00000103555 s 0.000001038775 s 1.00
sum / JaXPipe / tpu / BothRev 0.000001001225 s 9.997e-7 s 1.00
sum / Jax / tpu / BothRev 0.000001037075 s 0.0000010387 s 1.00
sum / HLOOpt / tpu / PreRev 0.000001016875 s 9.96275e-7 s 1.02
sum / HLOOpt / tpu / PostRev 0.00000104565 s 0.000001034625 s 1.01
sum / HLOOpt / tpu / BothRev 0.0000010146 s 9.887e-7 s 1.03
sum / PartOpt / tpu / PreRev 0.0000010358 s 0.00000104575 s 0.99
sum / PartOpt / tpu / PostRev 0.00000100605 s 9.97275e-7 s 1.01
sum / PartOpt / tpu / BothRev 0.0000010381 s 0.000001035775 s 1.00
sum / IPartOpt / tpu / PreRev 9.98175e-7 s 9.9675e-7 s 1.00
sum / IPartOpt / tpu / PostRev 0.0000010352500000000002 s 0.000001041975 s 0.99
sum / IPartOpt / tpu / BothRev 0.000001003325 s 9.8825e-7 s 1.02
sum / DefOpt / tpu / PreRev 0.00000103725 s 0.000001032125 s 1.00
sum / DefOpt / tpu / PostRev 0.0000010035 s 9.8955e-7 s 1.01
sum / DefOpt / tpu / BothRev 0.0000010392 s 0.0000010364 s 1.00
sum / IDefOpt / tpu / PreRev 0.0000010014 s 9.883e-7 s 1.01
sum / IDefOpt / tpu / PostRev 0.00000104105 s 0.000001033475 s 1.01
sum / IDefOpt / tpu / BothRev 0.000001004975 s 9.87125e-7 s 1.02
sum / JaXPipe / cpu / Primal 0.000014295 s 0.00000879642001564207 s 1.63
sum / Jax / cpu / Primal 0.000014231 s 0.000007896020024418249 s 1.80
sum / HLOOpt / cpu / Primal 0.000014448 s 0.00001246118001290597 s 1.16
sum / PartOpt / cpu / Primal 0.000013998 s 0.000008450039977105917 s 1.66
sum / IPartOpt / cpu / Primal 0.000014187 s 0.000008476499997414067 s 1.67
sum / DefOpt / cpu / Primal 0.000014199 s 0.000012441299986676311 s 1.14
sum / IDefOpt / cpu / Primal 0.000014264 s 0.000008302099977299805 s 1.72
sum / JaXPipe / cpu / Forward 0.000019493 s 0.000012332720007179888 s 1.58
sum / Jax / cpu / Forward 0.000020253 s 0.000012614460001714178 s 1.61
sum / HLOOpt / cpu / Forward 0.000019772 s 0.00001719820002108463 s 1.15
sum / PartOpt / cpu / Forward 0.000019661 s 0.000016864640010680886 s 1.17
sum / IPartOpt / cpu / Forward 0.000019903 s 0.000012690599969573669 s 1.57
sum / DefOpt / cpu / Forward 0.000019531 s 0.000017216219976035064 s 1.13
sum / IDefOpt / cpu / Forward 0.000019752 s 0.000012815299996873363 s 1.54
sum / JaXPipe / cpu / PreRev 0.000019029 s 0.00001255094000043755 s 1.52
sum / JaXPipe / cpu / PostRev 0.00001849 s 0.00001229197998327436 s 1.50
sum / JaXPipe / cpu / BothRev 0.000018473 s 0.00001570249994074402 s 1.18
sum / Jax / cpu / BothRev 0.000018731 s 0.00001219255997966684 s 1.54
sum / HLOOpt / cpu / PreRev 0.000018879 s 0.00001153561995124619 s 1.64
sum / HLOOpt / cpu / PostRev 0.00001923 s 0.00001561490000312915 s 1.23
sum / HLOOpt / cpu / BothRev 0.000018664 s 0.000013388559946179155 s 1.39
sum / PartOpt / cpu / PreRev 0.000018882 s 0.00001177519997327181 s 1.60
sum / PartOpt / cpu / PostRev 0.000018669 s 0.000011917339970750615 s 1.57
sum / PartOpt / cpu / BothRev 0.000019437 s 0.000011633939984676544 s 1.67
sum / IPartOpt / cpu / PreRev 0.000018952 s 0.000015914320001684246 s 1.19
sum / IPartOpt / cpu / PostRev 0.000019143 s 0.00001144826000199828 s 1.67
sum / IPartOpt / cpu / BothRev 0.000019054 s 0.000011945719979848943 s 1.60
sum / DefOpt / cpu / PreRev 0.000018739 s 0.0000116967799840495 s 1.60
sum / DefOpt / cpu / PostRev 0.000019374 s 0.000012046660003761644 s 1.61
sum / DefOpt / cpu / BothRev 0.00001911 s 0.00001183411999591044 s 1.61
sum / IDefOpt / cpu / PreRev 0.000018244 s 0.000011886359998243278 s 1.53
sum / IDefOpt / cpu / PostRev 0.000019606 s 0.000012097099952370627 s 1.62
sum / IDefOpt / cpu / BothRev 0.000018858 s 0.000011305200005153892 s 1.67
value_and_grad / JaXPipe / cpu / Primal 0.000016176940034711153 s 0.00001624343995899835 s 1.00
value_and_grad / Jax / cpu / Primal 0.000015957100013110905 s 0.00001564650000545953 s 1.02
value_and_grad / HLOOpt / cpu / Primal 0.00001549819998217572 s 0.000015223200025502592 s 1.02
value_and_grad / PartOpt / cpu / Primal 0.00001570629998241202 s 0.000015167499996096012 s 1.04
value_and_grad / IPartOpt / cpu / Primal 0.000015592080026181065 s 0.000015559860030407436 s 1.00
value_and_grad / DefOpt / cpu / Primal 0.000015601380009684363 s 0.00001555013997858623 s 1.00
value_and_grad / IDefOpt / cpu / Primal 0.000015423280037794028 s 0.000015306759996747134 s 1.01
value_and_grad / JaXPipe / cuda / Primal 0.000034144000000000004 s 0.000033759999999999995 s 1.01
value_and_grad / Jax / cuda / Primal 0.000034688 s 0.000034496 s 1.01
value_and_grad / HLOOpt / cuda / Primal 0.000033632 s 0.00003408 s 0.99
value_and_grad / PartOpt / cuda / Primal 0.000034112 s 0.000033728 s 1.01
value_and_grad / IPartOpt / cuda / Primal 0.000034304 s 0.000034048 s 1.01
value_and_grad / DefOpt / cuda / Primal 0.000035136000000000004 s 0.000033984 s 1.03
value_and_grad / IDefOpt / cuda / Primal 0.000035424 s 0.000033825 s 1.05
value_and_grad / JaXPipe / tpu / Primal 0 s 0 s 1
value_and_grad / Jax / tpu / Primal 0 s 0 s 1
value_and_grad / HLOOpt / tpu / Primal 0 s 0 s 1
value_and_grad / PartOpt / tpu / Primal 0 s 0 s 1
value_and_grad / IPartOpt / tpu / Primal 0 s 0 s 1
value_and_grad / DefOpt / tpu / Primal 0 s 0 s 1
value_and_grad / IDefOpt / tpu / Primal 0 s 0 s 1
value_and_grad / JaXPipe / cpu / Primal 0.000023285 s 0.00001624343995899835 s 1.43
value_and_grad / Jax / cpu / Primal 0.000022809 s 0.00001564650000545953 s 1.46
value_and_grad / HLOOpt / cpu / Primal 0.000022897 s 0.000015223200025502592 s 1.50
value_and_grad / PartOpt / cpu / Primal 0.000023071 s 0.000015167499996096012 s 1.52
value_and_grad / IPartOpt / cpu / Primal 0.000022859 s 0.000015559860030407436 s 1.47
value_and_grad / DefOpt / cpu / Primal 0.000022997 s 0.00001555013997858623 s 1.48
value_and_grad / IDefOpt / cpu / Primal 0.000022866 s 0.000015306759996747134 s 1.49
jaxmd20 / JaXPipe / cuda / Primal 0.001503109 s 0.001498085 s 1.00
jaxmd20 / Jax / cuda / Primal 0.0014648699999999 s 0.00143706 s 1.02
jaxmd20 / HLOOpt / cuda / Primal 0.00107338 s 0.001077793 s 1.00
jaxmd20 / PartOpt / cuda / Primal 0.001308517 s 0.0013310759999999 s 0.98
jaxmd20 / IPartOpt / cuda / Primal 0.001311557 s 0.001321539 s 0.99
jaxmd20 / DefOpt / cuda / Primal 0.000529634 s 0.000576802 s 0.92
jaxmd20 / IDefOpt / cuda / Primal 0.000523518 s 0.000512705 s 1.02
jaxmd20 / JaXPipe / cuda / Forward 0.0008669799999999 s 0.00085533 s 1.01
jaxmd20 / Jax / cuda / Forward 0.001823687 s 0.001797765 s 1.01
jaxmd20 / HLOOpt / cuda / Forward 0.000926339 s 0.000873794 s 1.06
jaxmd20 / PartOpt / cuda / Forward 0.000866948 s 0.000860515 s 1.01
jaxmd20 / IPartOpt / cuda / Forward 0.000868035 s 0.000874658 s 0.99
jaxmd20 / DefOpt / cuda / Forward 0.000889603 s 0.000864994 s 1.03
jaxmd20 / IDefOpt / cuda / Forward 0.0008664349999999 s 0.0008576979999999 s 1.01
jaxmd20 / JaXPipe / cuda / PreRev 0.0017435909999999 s 0.0017367399999999 s 1.00
jaxmd20 / JaXPipe / cuda / PostRev 0.005303476 s 0.005293839 s 1.00
jaxmd20 / JaXPipe / cuda / BothRev 0.001745127 s 0.001768356 s 0.99
jaxmd20 / Jax / cuda / BothRev 0.005270036 s 0.005272846 s 1.00
jaxmd20 / HLOOpt / cuda / PreRev 0.0017087749999999 s 0.001725541 s 0.99
jaxmd20 / HLOOpt / cuda / PostRev 0.005163316 s 0.005171214 s 1.00
jaxmd20 / HLOOpt / cuda / BothRev 0.001683015 s 0.001644709 s 1.02
jaxmd20 / PartOpt / cuda / PreRev 0.0018348549999999 s 0.0017877479999999 s 1.03
jaxmd20 / PartOpt / cuda / PostRev 0.0053552569999999 s 0.005330222 s 1.00
jaxmd20 / PartOpt / cuda / BothRev 0.001722758 s 0.001718148 s 1.00
jaxmd20 / IPartOpt / cuda / PreRev 0.001791111 s 0.001810405 s 0.99
jaxmd20 / IPartOpt / cuda / PostRev 0.005460468 s 0.00536619 s 1.02
jaxmd20 / IPartOpt / cuda / BothRev 0.0017080709999999 s 0.001705605 s 1.00
jaxmd20 / DefOpt / cuda / PreRev 0.0018017369999999 s 0.001828677 s 0.99
jaxmd20 / DefOpt / cuda / PostRev 0.002744171 s 0.002758728 s 0.99
jaxmd20 / DefOpt / cuda / BothRev 0.001725863 s 0.001713924 s 1.01
jaxmd20 / IDefOpt / cuda / PreRev 0.001764486 s 0.001793538 s 0.98
jaxmd20 / IDefOpt / cuda / PostRev 0.002209737 s 0.00220023 s 1.00
jaxmd20 / IDefOpt / cuda / BothRev 0.001731719 s 0.0017347889999999 s 1.00
jaxmd20 / JaXPipe / tpu / Primal 0.009282845625 s 0.0092661825 s 1.00
jaxmd20 / Jax / tpu / Primal 0.00927099875 s 0.009272955 s 1.00
jaxmd20 / HLOOpt / tpu / Primal 0.00916622125 s 0.009152614375 s 1.00
jaxmd20 / PartOpt / tpu / Primal 0.0091969375 s 0.009205384375 s 1.00
jaxmd20 / IPartOpt / tpu / Primal 0.009204061875 s 0.0092036 s 1.00
jaxmd20 / DefOpt / tpu / Primal 0.0087465562499999 s 0.008752913125 s 1.00
jaxmd20 / IDefOpt / tpu / Primal 0.00855528625 s 0.00855273375 s 1.00
jaxmd20 / JaXPipe / tpu / Forward 0.0170704025 s 0.01706695875 s 1.00
jaxmd20 / Jax / tpu / Forward 0.0187392649999999 s 0.018731186875 s 1.00
jaxmd20 / HLOOpt / tpu / Forward 0.017050126875 s 0.017046116875 s 1.00
jaxmd20 / PartOpt / tpu / Forward 0.0170731331249999 s 0.01706713125 s 1.00
jaxmd20 / IPartOpt / tpu / Forward 0.017066625 s 0.017071059375 s 1.00
jaxmd20 / DefOpt / tpu / Forward 0.017075628125 s 0.01707256125 s 1.00
jaxmd20 / IDefOpt / tpu / Forward 0.0170675025 s 0.01706959125 s 1.00
jaxmd20 / JaXPipe / tpu / PreRev 0.025011693125 s 0.0250080887499999 s 1.00
jaxmd20 / JaXPipe / tpu / PostRev 0.02187186875 s 0.021877951875 s 1.00
jaxmd20 / JaXPipe / tpu / BothRev 0.025026215 s 0.02500894375 s 1.00
jaxmd20 / Jax / tpu / BothRev 0.0218756606249999 s 0.02187837125 s 1.00
jaxmd20 / HLOOpt / tpu / PreRev 0.02501645 s 0.0249970018749999 s 1.00
jaxmd20 / HLOOpt / tpu / PostRev 0.0209730325 s 0.020705406875 s 1.01
jaxmd20 / HLOOpt / tpu / BothRev 0.024922951875 s 0.02490350875 s 1.00
jaxmd20 / PartOpt / tpu / PreRev 0.025019616875 s 0.02500404875 s 1.00
jaxmd20 / PartOpt / tpu / PostRev 0.0215276649999999 s 0.021508585625 s 1.00
jaxmd20 / PartOpt / tpu / BothRev 0.0249145106249999 s 0.024906398125 s 1.00
jaxmd20 / IPartOpt / tpu / PreRev 0.025023794375 s 0.025002875 s 1.00
jaxmd20 / IPartOpt / tpu / PostRev 0.02151866125 s 0.02151587 s 1.00
jaxmd20 / IPartOpt / tpu / BothRev 0.0249201981249999 s 0.02490401375 s 1.00
jaxmd20 / DefOpt / tpu / PreRev 0.0250178556249999 s 0.025004294375 s 1.00
jaxmd20 / DefOpt / tpu / PostRev 0.0187964875 s 0.01879673125 s 1.00
jaxmd20 / DefOpt / tpu / BothRev 0.0249161125 s 0.024908429375 s 1.00
jaxmd20 / IDefOpt / tpu / PreRev 0.0250211925 s 0.02499963125 s 1.00
jaxmd20 / IDefOpt / tpu / PostRev 0.01796021 s 0.0179288931249999 s 1.00
jaxmd20 / IDefOpt / tpu / BothRev 0.024921066875 s 0.02490134375 s 1.00
jaxmd40 / JaXPipe / cpu / Primal 0.073712136 s 0.078501014 s 0.94
jaxmd40 / Jax / cpu / Primal 0.058971163 s 0.06135705 s 0.96
jaxmd40 / HLOOpt / cpu / Primal 0.086776286 s 0.089335615 s 0.97
jaxmd40 / PartOpt / cpu / Primal 0.064984475 s 0.066812787 s 0.97
jaxmd40 / IPartOpt / cpu / Primal 0.058199601 s 0.0664021269999999 s 0.88
jaxmd40 / DefOpt / cpu / Primal 0.089436543 s 0.096352063 s 0.93
jaxmd40 / IDefOpt / cpu / Primal 0.083784705 s 0.094918816 s 0.88
jaxmd40 / JaXPipe / cpu / Forward 0.153904896 s 0.178721884 s 0.86
jaxmd40 / Jax / cpu / Forward 0.08231696 s 0.085444761 s 0.96
jaxmd40 / HLOOpt / cpu / Forward 0.153886849 s 0.172861772 s 0.89
jaxmd40 / PartOpt / cpu / Forward 0.1545946439999999 s 0.164854612 s 0.94
jaxmd40 / IPartOpt / cpu / Forward 0.14701933 s 0.175380794 s 0.84
jaxmd40 / DefOpt / cpu / Forward 0.151953942 s 0.171326453 s 0.89
jaxmd40 / IDefOpt / cpu / Forward 0.1477966879999999 s 0.162296659 s 0.91
jaxmd40 / JaXPipe / cpu / PreRev 0.2160650939999999 s 0.225926733 s 0.96
jaxmd40 / JaXPipe / cpu / PostRev 0.1362371359999999 s 0.140691395 s 0.97
jaxmd40 / JaXPipe / cpu / BothRev 0.213561996 s 0.226341833 s 0.94
jaxmd40 / Jax / cpu / BothRev 0.131048906 s 0.139427355 s 0.94
jaxmd40 / HLOOpt / cpu / PreRev 0.225377046 s 0.233465711 s 0.97
jaxmd40 / HLOOpt / cpu / PostRev 0.173531109 s 0.195115829 s 0.89
jaxmd40 / HLOOpt / cpu / BothRev 0.234282665 s 0.252705333 s 0.93
jaxmd40 / PartOpt / cpu / PreRev 0.226917946 s 0.229384134 s 0.99
jaxmd40 / PartOpt / cpu / PostRev 0.123175779 s 0.139846213 s 0.88
jaxmd40 / PartOpt / cpu / BothRev 0.238651269 s 0.262421069 s 0.91
jaxmd40 / IPartOpt / cpu / PreRev 0.230622181 s 0.2264402069999999 s 1.02
jaxmd40 / IPartOpt / cpu / PostRev 0.128587719 s 0.129780184 s 0.99
jaxmd40 / IPartOpt / cpu / BothRev 0.248020522 s 0.256815733 s 0.97
jaxmd40 / DefOpt / cpu / PreRev 0.215491352 s 0.2222731839999999 s 0.97
jaxmd40 / DefOpt / cpu / PostRev 0.175333665 s 0.1827078899999999 s 0.96
jaxmd40 / DefOpt / cpu / BothRev 0.24466169 s 0.258181016 s 0.95
jaxmd40 / IDefOpt / cpu / PreRev 0.223693448 s 0.213095599 s 1.05
jaxmd40 / IDefOpt / cpu / PostRev 0.170366415 s 0.176822343 s 0.96
jaxmd40 / IDefOpt / cpu / BothRev 0.229251341 s 0.27048417 s 0.85
jaxley_l5pc / JaXPipe / cuda / Primal 3.387076735496521 s
jaxley_l5pc / Jax / cuda / Primal 3.0272892609973496 s
jaxley_l5pc / HLOOpt / cuda / Primal 2.783555483500095 s
jaxley_l5pc / PartOpt / cuda / Primal 3.362598910000088 s
jaxley_l5pc / IPartOpt / cuda / Primal 3.3627612324999063 s
jaxley_l5pc / DefOpt / cuda / Primal 3.2315883549999853 s
jaxley_l5pc / IDefOpt / cuda / Primal 2.4416532440009178 s
jaxley_l5pc / NoScatterGatherOpts / cuda / Primal 3.39900577699882 s
jaxley_l5pc / JaXPipe / cuda / Forward 4.405809448500804 s
jaxley_l5pc / Jax / cuda / Forward 5.701747016497393 s
jaxley_l5pc / HLOOpt / cuda / Forward 4.514879172998917 s
jaxley_l5pc / PartOpt / cuda / Forward 4.4061166909996246 s
jaxley_l5pc / IPartOpt / cuda / Forward 4.40583299449645 s
jaxley_l5pc / DefOpt / cuda / Forward 4.410249729997304 s
jaxley_l5pc / IDefOpt / cuda / Forward 4.410543738002161 s
jaxley_l5pc / NoScatterGatherOpts / cuda / Forward 4.410188449499401 s
jaxley_l5pc / JaXPipe / tpu / Primal 45.38277049950011 s
jaxley_l5pc / Jax / tpu / Primal 30.25749352599996 s
jaxley_l5pc / HLOOpt / tpu / Primal 7.866905595999924 s
jaxley_l5pc / PartOpt / tpu / Primal 29.2853352969978 s
jaxley_l5pc / IPartOpt / tpu / Primal 29.28675943050075 s
jaxley_l5pc / DefOpt / tpu / Primal 33.17572148799991 s
jaxley_l5pc / IDefOpt / tpu / Primal 7.719404488001601 s
jaxley_l5pc / NoScatterGatherOpts / tpu / Primal 29.29138583849817 s
jaxley_l5pc / JaXPipe / tpu / Forward 14.423986808000336 s
jaxley_l5pc / Jax / tpu / Forward 60.27652003450021 s
jaxley_l5pc / HLOOpt / tpu / Forward 14.556906775498646 s
jaxley_l5pc / PartOpt / tpu / Forward 14.42279684299865 s
jaxley_l5pc / IPartOpt / tpu / Forward 14.422465295501752 s
jaxley_l5pc / DefOpt / tpu / Forward 14.427409824998904 s
jaxley_l5pc / IDefOpt / tpu / Forward 14.427485755000816 s
jaxley_l5pc / NoScatterGatherOpts / tpu / Forward 14.427506865002217 s
jaxley_l5pc / JaXPipe / cpu / Primal 1.511913058999653 s
jaxley_l5pc / Jax / cpu / Primal 1.0190414280004916 s
jaxley_l5pc / HLOOpt / cpu / Primal 0.8067630015002578 s
jaxley_l5pc / PartOpt / cpu / Primal 1.031235414999628 s
jaxley_l5pc / IPartOpt / cpu / Primal 1.0034160995001002 s
jaxley_l5pc / DefOpt / cpu / Primal 0.6768019065002591 s
jaxley_l5pc / IDefOpt / cpu / Primal 0.869091943499825 s
jaxley_l5pc / NoScatterGatherOpts / cpu / Primal 0.7466974170001777 s
jaxley_l5pc / JaXPipe / cpu / Forward 18.66026631350041 s
jaxley_l5pc / Jax / cpu / Forward 26.447296988999824 s
jaxley_l5pc / HLOOpt / cpu / Forward 18.76208826099992 s
jaxley_l5pc / PartOpt / cpu / Forward 18.82400553000025 s
jaxley_l5pc / IPartOpt / cpu / Forward 18.92444230199999 s
jaxley_l5pc / DefOpt / cpu / Forward 18.748458282999763 s
jaxley_l5pc / IDefOpt / cpu / Forward 18.595214767000016 s
jaxley_l5pc / NoScatterGatherOpts / cpu / Forward 18.55805930500037 s
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / JaXPipe / cuda / Primal 1.701457022 s 1.702424068 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / Jax / cuda / Primal 1.703768163 s 1.705451355 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / HLOOpt / cuda / Primal 1.715545969 s 1.715508605 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / PartOpt / cuda / Primal 1.695245413 s 1.6967465709999998 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IPartOpt / cuda / Primal 1.693645971 s 1.695191784 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / DefOpt / cuda / Primal 1.664803867 s 1.665287495 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IDefOpt / cuda / Primal 1.912197664 s 1.911164248 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / JaXPipe / tpu / Primal 4.005232194375 s 4.019771106875 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / Jax / tpu / Primal 3.038511031875 s 3.038608295 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / HLOOpt / tpu / Primal 3.12108057625 s 3.1210997818750004 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / PartOpt / tpu / Primal 3.05884849625 s 3.0589942700000003 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IPartOpt / tpu / Primal 3.058789881875 s 3.059039739375 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / DefOpt / tpu / Primal 2.102542103125 s 2.10257464625 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IDefOpt / tpu / Primal 4.355941740625 s 4.3560436725 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / JaXPipe / cpu / Primal 5.854209736 s 6.319441286 s 0.93
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / Jax / cpu / Primal 5.786406231 s 6.091046684 s 0.95
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / HLOOpt / cpu / Primal 5.774462324 s 6.052236802 s 0.95
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / PartOpt / cpu / Primal 5.846070446 s 6.141555782999999 s 0.95
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / IPartOpt / cpu / Primal 5.89435012 s 6.261373249 s 0.94
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / DefOpt / cpu / Primal 2.2922105580000003 s 2.439859331 s 0.94
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / IDefOpt / cpu / Primal 6.255501143 s 6.681259111999999 s 0.94

This comment was automatically generated by workflow using github-action-benchmark.

@avik-pal avik-pal force-pushed the ap/neuro_benchmark branch 2 times, most recently from 37aa1e7 to d70e93d Compare January 9, 2026 02:24
@avik-pal
Copy link
Collaborator Author

avik-pal commented Jan 9, 2026 

jaxley_l5pc             JaXPipe                 cpu             Forward         16.17755953    
/mnt2/avik-pal/.cache/bazel/_bazel_avik-pal/eca3efba73b0ff73b4ad16c26e74bbd7/sandbox/processwrapper-sandbox/16654/execroot/__main__/bazel-out/k8-dbg/bin/test/jaxley_test.runfiles/__main__/test/xprof_utils.py:62: UserWarning: xprof not found, falling back to timeit for profiling.
  warnings.warn("xprof not found, falling back to timeit for profiling.")
jaxley_l5pc             Jax                     cpu             Forward         20.42038421    
jaxley_l5pc             HLOOpt                  cpu             Forward         15.89919035    
jaxley_l5pc             PartOpt                 cpu             Forward         16.35003403    
jaxley_l5pc             IPartOpt                cpu             Forward         16.42652293    
jaxley_l5pc             DefOpt                  cpu             Forward         15.87073720    
jaxley_l5pc             IDefOpt                 cpu             Forward         15.76629268

~25% speedup no bad...

@avik-pal avik-pal force-pushed the ap/neuro_benchmark branch 2 times, most recently from cbe14c6 to 42571bd Compare January 9, 2026 21:12
@avik-pal avik-pal changed the base branch from main to ap/scatter_simplifications January 9, 2026 21:12
Base automatically changed from ap/scatter_simplifications to main January 10, 2026 00:51
@avik-pal avik-pal force-pushed the ap/neuro_benchmark branch 2 times, most recently from e158eb4 to d4c6fbc Compare January 12, 2026 13:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@github-actions github-actions[bot] github-actions[bot] left review comments

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@avik-pal